From be9fa192f06b515118ba8a8116ea43a4a45aa902 Mon Sep 17 00:00:00 2001 From: liangtaiwan Date: Fri, 18 Jan 2019 00:41:55 +0800 Subject: [PATCH] don't save if do not train --- examples/run_squad2.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/examples/run_squad2.py b/examples/run_squad2.py index fd35beef1e..558b24764e 100644 --- a/examples/run_squad2.py +++ b/examples/run_squad2.py @@ -1010,7 +1010,8 @@ def main(): # Save a trained model model_to_save = model.module if hasattr(model, 'module') else model # Only save the model it-self output_model_file = os.path.join(args.output_dir, "pytorch_model.bin") - torch.save(model_to_save.state_dict(), output_model_file) + if args.do_train: + torch.save(model_to_save.state_dict(), output_model_file) # Load a trained model that you have fine-tuned model_state_dict = torch.load(output_model_file)