diff --git a/justfile b/justfile
index fe758bb..e421379 100644
--- a/justfile
+++ b/justfile
@@ -5,6 +5,7 @@ build:
 build-cuda:
     cd pyr && cargo build --release --features cuda
 train level:
+    python -m mlgame -f 400 -i ./ml/ml_play_pyr_test.py -i ./ml/ml_play_manual_2P.py . --sound off --level {{level}} --game_times 3 | grep Rewards
 run level:
     python -m mlgame -f 400 -i ./ml/ml_play_pyr_test.py -i ./ml/ml_play_manual_2P.py . --sound off --level {{level}} --game_times 3
 clean:
diff --git a/pyr/src/app/agent.rs b/pyr/src/app/agent.rs
index b28efb0..4c23204 100644
--- a/pyr/src/app/agent.rs
+++ b/pyr/src/app/agent.rs
@@ -186,7 +186,6 @@ impl AIAgent {
         // Compare the estimated rewards with the maximum expected rewards and
         // perform the backward step.
         let loss = huber_loss(1.0_f32)(&x, &y);
-        log::trace!("loss: {:?}", loss);
         self.optimizer
             .backward_step(&Tensor::new(&[loss], &self.device).unwrap())
             .unwrap();
@@ -195,6 +194,8 @@ impl AIAgent {
         self.memory.clear();
         if CONFIG.train {
             self.var_map.save("model.bin").unwrap();
+            log::info!("Rewards {}", self.accumulate_rewards as i64);
+            self.accumulate_rewards = 0.0;
             log::info!("model.bin saved!");
         }
     }