Fix yaml load (#5597)

intel-analytics · Aug 31, 2022 · b18aec8 · b18aec8
1 parent 53dfd78
commit b18aec8
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/python/friesian/example/dien/e2e_infer.py b/python/friesian/example/dien/e2e_infer.py
@@ -325,7 +325,7 @@ def process_single_meta(row):
     df = full_tbl.df.repartition(num_tasks)
     df = df.sortWithinPartitions("item_hist_seq_len", ascending=False)
 
-    config = yaml.load(open("config_runtime.yaml", "r"), Loader=yaml.FullLoader)
+    config = yaml.safe_load(open("config_runtime.yaml", "r"))
     rdd = df.rdd
     eval_res = rdd.mapPartitions(lambda iter: infer_main(iter, config)).collect()
     data_size = 0