winglian commited on
Commit
915c56c
1 Parent(s): cd23959

Update finetune.py

Browse files
Files changed (1) hide show
  1. scripts/finetune.py +4 -5
scripts/finetune.py CHANGED
@@ -229,11 +229,10 @@ def train(
229
  f"Training Completed!!! Saving pre-trained model to {cfg.output_dir}"
230
  )
231
 
232
- if cfg.adapter == 'lora':
233
- trainer.save_pretrained(cfg.output_dir)
234
- else:
235
- # TODO do we need this fix? https://huggingface.co/docs/accelerate/usage_guides/fsdp#saving-and-loading
236
- trainer.save_model(cfg.output_dir)
237
 
238
 
239
  if __name__ == "__main__":
 
229
  f"Training Completed!!! Saving pre-trained model to {cfg.output_dir}"
230
  )
231
 
232
+
233
+ # TODO do we need this fix? https://huggingface.co/docs/accelerate/usage_guides/fsdp#saving-and-loading
234
+ trainer.save_pretrained(cfg.output_dir)
235
+ # trainer.save_model(cfg.output_dir) # TODO this may be needed for deepspeed to work? need to review another time
 
236
 
237
 
238
  if __name__ == "__main__":