From b18aec88c05167277560d01a2728e8016189274a Mon Sep 17 00:00:00 2001 From: Kai Huang Date: Wed, 31 Aug 2022 16:46:05 +0800 Subject: [PATCH] Fix yaml load (#5597) --- python/friesian/example/dien/e2e_infer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/friesian/example/dien/e2e_infer.py b/python/friesian/example/dien/e2e_infer.py index c3d18187a1c..cde3f31b214 100644 --- 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