fix a bug in mnist_ddp example

Azure · Dec 13, 2023 · 42ac3c2 · 42ac3c2
1 parent b1434fb
commit 42ac3c2
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/examples/mnist_ddp.py b/examples/mnist_ddp.py
@@ -170,7 +170,9 @@ def main():
 
     if args.local_rank == 0:
         dataset1 = datasets.MNIST('./data', train=True, download=True, transform=transform)
-    torch.cuda.synchronize()
+
+    torch.distributed.barrier()
+
     if args.local_rank > 0:
         dataset1 = datasets.MNIST('./data', train=True, download=False, transform=transform)