Merge pull request #693 from SamuelGabriel/patch-1

Let only the _globally_ 0th rank write checkpoints in `train.py`
3 years ago · b79dfd4fc2
parent b7a568f065 7c19c35d9f
commit b79dfd4fc2
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: