This commit is contained in:
Julien Chaumond
2019-11-05 19:06:12 -05:00
parent de890ae67d
commit 30968d70af
6 changed files with 27 additions and 17 deletions

View File

@@ -237,7 +237,7 @@ def main():
# Save a trained model
if args.do_train:
# Save a trained model, configuration and tokenizer
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 itself
# If we save using the predefined names, we can load using `from_pretrained`
output_model_file = os.path.join(args.output_dir, WEIGHTS_NAME)