Revert "Fix: save model/model.module"

This reverts commit 00e9c4cc96.
This commit is contained in:
wangfei
2019-08-18 11:03:12 +08:00
parent 00e9c4cc96
commit 1ef41b8337
2 changed files with 6 additions and 8 deletions

View File

@@ -322,8 +322,7 @@ def main():
# Save a trained model
if args.local_rank == -1 or torch.distributed.get_rank() == 0:
logging.info("** ** * Saving fine-tuned model ** ** * ")
model_to_save = model.module if hasattr(model, 'module') else model # Take care of distributed/parallel training
model_to_save.save_pretrained(args.output_dir)
model.save_pretrained(args.output_dir)
tokenizer.save_pretrained(args.output_dir)

View File

@@ -610,8 +610,7 @@ def main():
# Save a trained model
if args.do_train and (args.local_rank == -1 or torch.distributed.get_rank() == 0):
logger.info("** ** * Saving fine - tuned model ** ** * ")
model_to_save = model.module if hasattr(model, 'module') else model # Take care of distributed/parallel training
model_to_save.save_pretrained(args.output_dir)
model.save_pretrained(args.output_dir)
tokenizer.save_pretrained(args.output_dir)