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:
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user