SUMIN/HuggingFace_transformer
1
0
Fork 0
Code Issues Pull Requests Actions 7 Packages Projects Releases Wiki Activity

don't save if do not train

Browse Source
This commit is contained in:
liangtaiwan
2019-01-18 00:41:55 +08:00
parent 647c983530
commit be9fa192f0
1 changed files with 2 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
examples/run_squad2.py
View File

@@ -1010,6 +1010,7 @@ def main():
# Save a trained model # Save a trained model
model_to_save = model.module if hasattr(model, 'module') else model # Only save the model it-self 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") 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) torch.save(model_to_save.state_dict(), output_model_file)
# Load a trained model that you have fine-tuned # Load a trained model that you have fine-tuned