fix printing of loading cpt

ai-forever · Aug 12, 2021 · 4eb1fbe · 4eb1fbe
1 parent 25719f1
commit 4eb1fbe
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/src/utils.py b/src/utils.py
@@ -329,10 +329,6 @@ def load_checkpoint(model, optimizer, lr_scheduler, args, deepspeed=False):
         # Checkpoint.
         checkpoint_name = get_checkpoint_name(args.load, iteration, release)
 
-        if mpu.get_data_parallel_rank() == 0:
-            print('global rank {} is loading checkpoint {}'.format(
-                torch.distributed.get_rank(), checkpoint_name))
-
         # Load the checkpoint.
         if os.path.isfile(checkpoint_name):
             sd = torch.load(checkpoint_name, map_location='cpu')
@@ -344,6 +340,10 @@ def load_checkpoint(model, optimizer, lr_scheduler, args, deepspeed=False):
             )
             sd = torch.load(checkpoint_name, map_location='cpu')
 
+        if mpu.get_data_parallel_rank() == 0:
+            print('global rank {} is loading checkpoint {}'.format(
+                torch.distributed.get_rank(), checkpoint_name))
+
         if isinstance(model, torchDDP):
             model = model.module