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.<field>.len() for all memcpy_dtoh host
allocations to prevent assert!(dst.len() >= src.len()) panics.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-04 01:07:58 +02:00
parent e4dc86bf85
commit 0e6a3c4485

View File

@@ -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