Global instead of local rank.

4 years ago · 7c19c35d9f
parent b7a568f065
commit 7c19c35d9f
1 changed files with 1 additions and 1 deletions
--- 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: