Update finetune.py
This commit is contained in:
@@ -229,11 +229,10 @@ def train(
|
|||||||
f"Training Completed!!! Saving pre-trained model to {cfg.output_dir}"
|
f"Training Completed!!! Saving pre-trained model to {cfg.output_dir}"
|
||||||
)
|
)
|
||||||
|
|
||||||
if cfg.adapter == 'lora':
|
|
||||||
trainer.save_pretrained(cfg.output_dir)
|
# TODO do we need this fix? https://huggingface.co/docs/accelerate/usage_guides/fsdp#saving-and-loading
|
||||||
else:
|
trainer.save_pretrained(cfg.output_dir)
|
||||||
# TODO do we need this fix? https://huggingface.co/docs/accelerate/usage_guides/fsdp#saving-and-loading
|
# trainer.save_model(cfg.output_dir) # TODO this may be needed for deepspeed to work? need to review another time
|
||||||
trainer.save_model(cfg.output_dir)
|
|
||||||
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
if __name__ == "__main__":
|
||||||
|
|||||||
Reference in New Issue
Block a user