fixed save race condition

0399d32c · Raul Puri · GitHub · 3573423f · 0399d32c
Unverified Commit 0399d32c authored May 14, 2019 by Raul Puri Committed by GitHub May 14, 2019
Hide whitespace changes
Inline Side-by-side

Showing with 1 addition and 1 deletion

utils.py utils.py +1 -1

No files found.
--- a/utils.py
+++ b/utils.py
@@ -149,7 +149,7 @@ def save_checkpoint(model_suffix, epoch, i, model, optimizer, lr_scheduler, args
                 np.random.get_state(),
                 random.getstate())
    if not (torch.distributed.is_initialized() and \
-            torch.distributed.get_rank() > 1):
+            torch.distributed.get_rank() > 0):
        if not os.path.exists(checkpoint_dir):
            os.makedirs(checkpoint_dir)
        total_iters = args.train_iters * (epoch-1) + i