From 2d73591a1831e80d0743b514d7f0138c4879e37b Mon Sep 17 00:00:00 2001 From: Bilal Khan Date: Wed, 27 Nov 2019 19:13:10 -0600 Subject: [PATCH] Stop saving current epoch --- examples/run_lm_finetuning.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/examples/run_lm_finetuning.py b/examples/run_lm_finetuning.py index 9bdbf9ca56..5e7683b85d 100644 --- a/examples/run_lm_finetuning.py +++ b/examples/run_lm_finetuning.py @@ -290,8 +290,7 @@ def train(args, train_dataset, model, tokenizer): torch.save(optimizer.state_dict(), os.path.join(output_dir, 'optimizer.pt')) torch.save(scheduler.state_dict(), os.path.join(output_dir, 'scheduler.pt')) - torch.save(epoch, os.path.join(output_dir, 'training_state.pt')) - logger.info("Saving training state to %s", output_dir) + logger.info("Saving optimizer and scheduler states to %s", output_dir) if args.max_steps > 0 and global_step > args.max_steps: epoch_iterator.close()