From 1ef0f89a0c0125eb189cb0932d8da78919ca17d5 Mon Sep 17 00:00:00 2001
From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com>
Date: Wed, 22 Feb 2023 00:14:24 +0530
Subject: [PATCH] add util for getting the base model

---
 src/peft/peft_model.py | 11 +++++++----
 1 file changed, 7 insertions(+), 4 deletions(-)

diff --git a/src/peft/peft_model.py b/src/peft/peft_model.py
index f12c5be..e20e64d 100644
--- a/src/peft/peft_model.py
+++ b/src/peft/peft_model.py
@@ -284,10 +284,7 @@ class PeftModel(PushToHubMixin, torch.nn.Module):
         """
         Forward pass of the model.
         """
-        if isinstance(self.peft_config, PromptLearningConfig):
-            return self.base_model(*args, **kwargs)
-        else:
-            return self.base_model.model(*args, **kwargs)
+        return self.get_base_model()(*args, **kwargs)
 
     @contextmanager
     def disable_adapter(self):
@@ -305,6 +302,12 @@ class PeftModel(PushToHubMixin, torch.nn.Module):
         else:
             self.base_model.enable_adapter_layers()
 
+    def get_base_model(self):
+        """
+        Returns the base model.
+        """
+        return self.base_model if isinstance(self.peft_config, PromptLearningConfig) else self.base_model.model
+
 
 class PeftModelForSequenceClassification(PeftModel):
     """