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