Use model.save_pretrained() to save model at the end of training

yukw777 · Apr 5, 2023 · 51ca51d · 51ca51d
1 parent c7164ee
commit 51ca51d
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/train.py b/train.py
@@ -165,7 +165,7 @@ def train() -> None:
  ),
  )
  trainer.train()
- trainer.save_model()
+ model.save_pretrained(training_args.output_dir)
 
 
 if __name__ == "__main__":