Skip to content

GitLab

Commit 123da5a2 authored by yzy5630's avatar yzy5630
Browse files

fix errors for lm_finetuning examples

parent 60a1bdcd
Hide whitespace changes
Inline Side-by-side
Showing with 8 additions and 7 deletions
examples/lm_finetuning/finetune_on_pregenerated.py
View file @ 123da5a2
...@@ -325,15 +325,16 @@ def main(): ...@@ -325,15 +325,16 @@ def main():
global_step += 1 global_step += 1
# Save a trained model # Save a trained model
logging.info("** ** * Saving fine-tuned model ** ** * ") if torch.distributed.get_rank() == 0:
model_to_save = model.module if hasattr(model, 'module') else model # Only save the model it-self logging.info("** ** * Saving fine-tuned 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, WEIGHTS_NAME) output_model_file = os.path.join(args.output_dir, WEIGHTS_NAME)
output_config_file = os.path.join(args.output_dir, CONFIG_NAME) output_config_file = os.path.join(args.output_dir, CONFIG_NAME)
torch.save(model_to_save.state_dict(), output_model_file) torch.save(model_to_save.state_dict(), output_model_file)
model_to_save.config.to_json_file(output_config_file) model_to_save.config.to_json_file(output_config_file)
tokenizer.save_vocabulary(args.output_dir) tokenizer.save_vocabulary(args.output_dir)
if __name__ == '__main__': if __name__ == '__main__':
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment