Update examples/impala/impala_multi_node_ray.py

Co-authored-by: Vincent Moens <vincentmoens@gmail.com>
pytorch · Nov 22, 2023 · 4a061b5 · 4a061b5
1 parent 72a3c6e
commit 4a061b5
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/examples/impala/impala_multi_node_ray.py b/examples/impala/impala_multi_node_ray.py
@@ -182,7 +182,7 @@ def main(cfg: "DictConfig"):  # noqa: F821
 
             # Create a single batch of trajectories
             stacked_data = torch.stack(accumulator, dim=0)
-            stacked_data = stacked_data.to(device)
+            stacked_data = stacked_data.to(device, non_blocking=True)
 
             # Compute advantage
             stacked_data = adv_module(stacked_data)