don't save if do not train

2019-01-18 00:41:55 +08:00
parent 647c983530
commit be9fa192f0
1 changed files with 2 additions and 1 deletions
--- a/examples/run_squad2.py
+++ b/examples/run_squad2.py
@@ -1010,6 +1010,7 @@ 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")
+    if args.do_train:
        torch.save(model_to_save.state_dict(), output_model_file)

    # Load a trained model that you have fine-tuned