Merge pull request #27 from NanoCode012/patch-1
Browse files- scripts/finetune.py +1 -1
scripts/finetune.py
CHANGED
|
@@ -236,7 +236,7 @@ def train(
|
|
| 236 |
logging.info(f"Training Completed!!! Saving pre-trained model to {cfg.output_dir}")
|
| 237 |
|
| 238 |
# TODO do we need this fix? https://huggingface.co/docs/accelerate/usage_guides/fsdp#saving-and-loading
|
| 239 |
-
|
| 240 |
# trainer.save_model(cfg.output_dir) # TODO this may be needed for deepspeed to work? need to review another time
|
| 241 |
|
| 242 |
|
|
|
|
| 236 |
logging.info(f"Training Completed!!! Saving pre-trained model to {cfg.output_dir}")
|
| 237 |
|
| 238 |
# TODO do we need this fix? https://huggingface.co/docs/accelerate/usage_guides/fsdp#saving-and-loading
|
| 239 |
+
model.save_pretrained(cfg.output_dir)
|
| 240 |
# trainer.save_model(cfg.output_dir) # TODO this may be needed for deepspeed to work? need to review another time
|
| 241 |
|
| 242 |
|