From af7414a67d5aecc68b59e6cee87f55dd05b9b854 Mon Sep 17 00:00:00 2001
From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com>
Date: Wed, 15 Feb 2023 15:44:26 +0530
Subject: [PATCH] fix forward signature

---
 src/peft/peft_model.py  | 1 +
 src/peft/utils/other.py | 2 +-
 2 files changed, 2 insertions(+), 1 deletion(-)

diff --git a/src/peft/peft_model.py b/src/peft/peft_model.py
index e2932d9..cca1c02 100644
--- a/src/peft/peft_model.py
+++ b/src/peft/peft_model.py
@@ -81,6 +81,7 @@ class PeftModel(PushToHubMixin, torch.nn.Module):
             self.modules_to_save = self.peft_config.modules_to_save
             _set_trainable(self)
         self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.forward.__func__.__signature__ = inspect.signature(self.base_model.forward)
 
     def save_pretrained(self, save_directory, **kwargs):
         r"""
diff --git a/src/peft/utils/other.py b/src/peft/utils/other.py
index ac93a1e..82a4a45 100644
--- a/src/peft/utils/other.py
+++ b/src/peft/utils/other.py
@@ -42,7 +42,7 @@ def prepare_model_for_training(model, output_embedding_layer_name="lm_head"):
     """
     loaded_in_8bit = getattr(model, "is_loaded_in_8bit", False)
 
-    for name,param in model.named_parameters():
+    for name, param in model.named_parameters():
         # freeze base model's layers
         param.requires_grad = False