Global instead of local rank.

SamuelGabriel · web-flow · commit 7c19c35d9f11 · 2021-06-09T19:11:58.000+02:00
diff --git a/train.py b/train.py
@@ -561,7 +561,7 @@ def main():
     best_epoch = None
     saver = None
     output_dir = None
-    if args.local_rank == 0:
+    if args.rank == 0:
         if args.experiment:
             exp_name = args.experiment
         else:
