Update examples/impala/impala_multi_node_ray.py

Co-authored-by: Vincent Moens <vincentmoens@gmail.com>
pytorch · Nov 22, 2023 · e7069e4 · e7069e4
1 parent 4a061b5
commit e7069e4
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
@@ -203,7 +203,7 @@ def main(cfg: "DictConfig"):  # noqa: F821
                 num_network_updates += 1
 
                 # Get a data batch
-                batch = batch.to(device)
+                batch = batch.to(device, non_blocking=True)
 
                 # Forward pass loss
                 loss = loss_module(batch)