diff --git a/crates/ml/src/trainers/dqn/fused_training.rs b/crates/ml/src/trainers/dqn/fused_training.rs index bf7018dd7..8d06c0a19 100644 --- a/crates/ml/src/trainers/dqn/fused_training.rs +++ b/crates/ml/src/trainers/dqn/fused_training.rs @@ -792,7 +792,10 @@ impl FusedTrainingCtx { &self.stream, ).map_err(|e| anyhow::anyhow!("PER seg_tree_update: {e}"))?; - eprintln!("FUSED_DIAG: step {} — PER update done", self.steps_since_varmap_sync); + eprintln!("FUSED_DIAG: step {} — PER update done, syncing full step", self.steps_since_varmap_sync); + self.trainer.stream().synchronize() + .map_err(|e| anyhow::anyhow!("post-step sync: {e}"))?; + eprintln!("FUSED_DIAG: step {} — full step GPU sync complete", self.steps_since_varmap_sync); // Bookkeeping. agent.fused_post_step_gpu_bookkeeping() .map_err(|e| anyhow::anyhow!("Fused GPU bookkeeping: {e}"))?;