fixed exit interval

c380000e · Mostofa Patwary · 097fa65b · c380000e
Commit c380000e authored Feb 03, 2021 by Mostofa Patwary
Hide whitespace changes
Inline Side-by-side

Showing with 3 additions and 3 deletions

megatron/training.py megatron/training.py +3 -3

No files found.
--- a/megatron/training.py
+++ b/megatron/training.py
@@ -884,9 +884,9 @@ def train(forward_step_func, model, optimizer, lr_scheduler,
        # Exiting based on iterations        
        if args.exit_interval and iteration % args.exit_interval == 0:
-            #if not saved_checkpoint:
+            if not saved_checkpoint:
-            #    save_checkpoint_and_time(iteration, model, optimizer,
+                save_checkpoint_and_time(iteration, model, optimizer,
-            #                             lr_scheduler)
+                                         lr_scheduler)
            torch.distributed.barrier()
            print_datetime('exiting program at iteration {}'.format(iteration))                
            sys.exit()