From 915c56cd97bc9c336bb354b457a43118ca527642 Mon Sep 17 00:00:00 2001 From: Wing Lian Date: Tue, 9 May 2023 15:05:39 -0400 Subject: [PATCH] Update finetune.py --- scripts/finetune.py | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/scripts/finetune.py b/scripts/finetune.py index c05023f40..915ba1de1 100644 --- a/scripts/finetune.py +++ b/scripts/finetune.py @@ -229,11 +229,10 @@ def train( f"Training Completed!!! Saving pre-trained model to {cfg.output_dir}" ) - if cfg.adapter == 'lora': - trainer.save_pretrained(cfg.output_dir) - else: - # TODO do we need this fix? https://huggingface.co/docs/accelerate/usage_guides/fsdp#saving-and-loading - trainer.save_model(cfg.output_dir) + + # TODO do we need this fix? https://huggingface.co/docs/accelerate/usage_guides/fsdp#saving-and-loading + trainer.save_pretrained(cfg.output_dir) + # trainer.save_model(cfg.output_dir) # TODO this may be needed for deepspeed to work? need to review another time if __name__ == "__main__":