Commit c380000e authored by Mostofa Patwary's avatar Mostofa Patwary
Browse files

fixed exit interval

parent 097fa65b

megatron/training.py

+3 −3
Original line number Diff line number Diff line
@@ -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:

            #    save_checkpoint_and_time(iteration, model, optimizer,

            #                             lr_scheduler)

            if not saved_checkpoint:

                save_checkpoint_and_time(iteration, model, optimizer,

                                         lr_scheduler)

            torch.distributed.barrier()

            print_datetime('exiting program at iteration {}'.format(iteration))                

            sys.exit()