diff --git a/rdagent/scenarios/data_mining/experiment/model_template/train.py b/rdagent/scenarios/data_mining/experiment/model_template/train.py
index 26b343ff..f00ade8e 100644
--- a/rdagent/scenarios/data_mining/experiment/model_template/train.py
+++ b/rdagent/scenarios/data_mining/experiment/model_template/train.py
@@ -75,9 +75,21 @@ model = model_cls(num_features=num_features, num_timesteps=num_timesteps).to(dev
 optimizer = torch.optim.Adam(model.parameters(), lr=0.0001)
 criterion = nn.CrossEntropyLoss()
 
-# Train the model
 
-for i in range(10):
+# Train the model
+def eval_auc(model):
+    y_pred = []
+    for data in test_dataloader:
+        x, y = data
+        out = model(x)
+        y_pred.append(out.cpu().detach().numpy())
+    return roc_auc_score(y_test, np.concatenate(y_pred))
+
+
+best = 0.0
+best_model = None
+
+for i in range(15):
     for data in train_dataloader:
         x, y = data
         out = model(x)
@@ -85,11 +97,15 @@ for i in range(10):
         loss = criterion(out.squeeze(), y)
         loss.backward()
         optimizer.step()
+    roc = eval_auc(model)
+    if roc > best:
+        best = roc
+        best_model = model
 
 y_pred = []
 for data in test_dataloader:
     x, y = data
-    out = model(x)
+    out = best_model(x)
     y_pred.append(out.cpu().detach().numpy())
 
 acc = roc_auc_score(y_test, np.concatenate(y_pred))
diff --git a/rdagent/scenarios/data_mining/experiment/prompts.yaml b/rdagent/scenarios/data_mining/experiment/prompts.yaml
index 50e0129f..43db687f 100644
--- a/rdagent/scenarios/data_mining/experiment/prompts.yaml
+++ b/rdagent/scenarios/data_mining/experiment/prompts.yaml
@@ -34,7 +34,11 @@ dm_model_interface: |-
   ```
   No other parameters will be passed to the model so give other parameters a default value or just make them static.
 
-  Remember to permute the input tensor since the input tensor is in the shape of (batch_size, num_features, num_timesteps) and a normal time series model is expecting the input tensor in the shape of (batch_size, num_timesteps, num_features).
+  The input tensor shape is (batch_size, num_features, num_timesteps) which is different from the normal time series input shape of (batch_size, num_timesteps, num_features). Please write code accordingly. 
+  
+  Note that for nn.Conv1d() layers, please do not permute the input tensor as the in_channel dimension should match the num_feature dimension.
+
+  The output shape should be (batch_size, 1) with sigmoid activation since we have binary labels. 
 
   Don't write any try-except block in your python code. The user will catch the exception message and provide the feedback to you. Also, don't write main function in your python code. The user will call the forward method in the model_cls to get the output tensor.