From b8a699a0df622a76a9cd944f3ed2b686e51e9cb8 Mon Sep 17 00:00:00 2001
From: Yuante Li <104308117+WinstonLiyt@users.noreply.github.com>
Date: Mon, 17 Mar 2025 21:43:26 +0800
Subject: [PATCH] fix: fix a minor bug (#694)

---
 .../coder/data_science/raw_data_loader/prompts.yaml           | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/rdagent/components/coder/data_science/raw_data_loader/prompts.yaml b/rdagent/components/coder/data_science/raw_data_loader/prompts.yaml
index 6f97d407..7be6340f 100644
--- a/rdagent/components/coder/data_science/raw_data_loader/prompts.yaml
+++ b/rdagent/components/coder/data_science/raw_data_loader/prompts.yaml
@@ -294,8 +294,8 @@ spec:
             val_y=val_y,
             test_X=X_test_transformed
           )
-          val_preds_dict[model_name] = val_pred
-          test_preds_dict[model_name] = test_pred
+          val_preds_dict[model_module] = val_pred
+          test_preds_dict[model_module] = test_pred
           {% endfor %}
           final_pred = ensemble_workflow(test_preds_dict, val_preds_dict, val_y)
           {% endraw %}