diff --git a/crates/ml/tests/smoke_test_real_data.rs b/crates/ml/tests/smoke_test_real_data.rs
index f66d32a20..1c842aa34 100644
--- a/crates/ml/tests/smoke_test_real_data.rs
+++ b/crates/ml/tests/smoke_test_real_data.rs
@@ -569,9 +569,9 @@ mod gpu_smoke {
         assert_eq!(batch.n_episodes * batch.timesteps, expected_total);
 
         // Download actions and rewards from GPU to host for validation
-        let mut actions_host = vec![0i32; expected_total];
+        let mut actions_host = vec![0i32; batch.actions.len()];
         stream.memcpy_dtoh(&batch.actions, &mut actions_host).unwrap();
-        let mut rewards_host = vec![0.0f32; expected_total];
+        let mut rewards_host = vec![0.0f32; batch.rewards.len()];
         stream.memcpy_dtoh(&batch.rewards, &mut rewards_host).unwrap();
 
         // Branching DQN: 9×3×3 = 81 factored actions (0..81)
@@ -649,7 +649,7 @@ mod gpu_smoke {
         assert_eq!(batch.n_episodes * batch.timesteps, expected_total);
 
         // Download rewards from GPU to host for validation
-        let mut rewards_host = vec![0.0f32; expected_total];
+        let mut rewards_host = vec![0.0f32; batch.rewards.len()];
         stream.memcpy_dtoh(&batch.rewards, &mut rewards_host).unwrap();
 
         for (i, &r) in rewards_host.iter().enumerate() {
@@ -730,9 +730,9 @@ mod gpu_smoke {
         assert_eq!(batch.n_episodes * batch.timesteps, expected_total);
 
         // Download actions and rewards from GPU to host for validation
-        let mut actions_host = vec![0i32; expected_total];
+        let mut actions_host = vec![0i32; batch.actions.len()];
         stream.memcpy_dtoh(&batch.actions, &mut actions_host).unwrap();
-        let mut rewards_host = vec![0.0f32; expected_total];
+        let mut rewards_host = vec![0.0f32; batch.rewards.len()];
         stream.memcpy_dtoh(&batch.rewards, &mut rewards_host).unwrap();
 
         // Rewards should be finite