Fix single-GPU fallback for examples/distributed

4a8cf7ad · mcarilli · GitHub · 31931985 · 4a8cf7ad
Unverified Commit 4a8cf7ad authored Jul 09, 2018 by mcarilli Committed by GitHub Jul 09, 2018
Hide whitespace changes
Inline Side-by-side

Showing with 2 additions and 1 deletion

examples/distributed/main.py examples/distributed/main.py +2 -1

No files found.
--- a/examples/distributed/main.py
+++ b/examples/distributed/main.py
@@ -201,7 +201,8 @@ def test():
 for epoch in range(1, args.epochs + 1):
    #=====START: ADDED FOR DISTRIBUTED======
-    train_sampler.set_epoch(epoch)
+    if args.distributed:
+        train_sampler.set_epoch(epoch)
    #=====END:   ADDED FOR DISTRIBUTED======
    train(epoch)