From b1af297707a72ef8d2234d865f5d338f8d94f95f Mon Sep 17 00:00:00 2001 From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> Date: Fri, 10 Feb 2023 00:06:25 +0530 Subject: [PATCH] make `save_pretrained` work in a way training could be resumed --- src/peft/peft_model.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/src/peft/peft_model.py b/src/peft/peft_model.py index 49bfce1..2db3e76 100644 --- a/src/peft/peft_model.py +++ b/src/peft/peft_model.py @@ -98,9 +98,6 @@ class PeftModel(PushToHubMixin, torch.nn.Module): raise ValueError(f"Provided path ({save_directory}) should be a directory, not a file") os.makedirs(save_directory, exist_ok=True) - for param in self.parameters(): - param.requires_grad = False # freeze the model - # save only the trainable weights output_state_dict = get_peft_model_state_dict(self, kwargs.get("state_dict", None)) torch.save(output_state_dict, os.path.join(save_directory, WEIGHTS_NAME)) @@ -112,8 +109,10 @@ class PeftModel(PushToHubMixin, torch.nn.Module): if isinstance(self.peft_config, PromptLearningConfig) else self.base_model.model.__dict__.get("name_or_path", None) ) + inference_mode = self.peft_config.inference_mode self.peft_config.inference_mode = True self.peft_config.save_pretrained(save_directory) + self.peft_config.inference_mode = inference_mode @classmethod def from_pretrained(cls, model, model_id, **kwargs):