From 0e6a3c4485be4471d3fa4d8d18fe6c7f7384b404 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Sat, 4 Apr 2026 01:07:58 +0200 Subject: [PATCH] fix: size memcpy host vecs from GPU buffer length (not expected_total) GPU buffers may be larger than n_episodes * timesteps due to padding/alignment; use batch..len() for all memcpy_dtoh host allocations to prevent assert!(dst.len() >= src.len()) panics. Co-Authored-By: Claude Sonnet 4.6 --- crates/ml/tests/smoke_test_real_data.rs | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) 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