From 18ccde8e86a1dcebc1d2ddf85350d09a341342d2 Mon Sep 17 00:00:00 2001
From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com>
Date: Tue, 4 Apr 2023 19:44:58 +0530
Subject: [PATCH] =?UTF-8?q?fixing=20=F0=9F=90=9B?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 src/peft/peft_model.py          | 6 ++++--
 src/peft/tuners/lora.py         | 2 +-
 src/peft/utils/save_and_load.py | 4 ++--
 3 files changed, 7 insertions(+), 5 deletions(-)

diff --git a/src/peft/peft_model.py b/src/peft/peft_model.py
index b98fc9f..1c3e7e1 100644
--- a/src/peft/peft_model.py
+++ b/src/peft/peft_model.py
@@ -112,7 +112,9 @@ class PeftModel(PushToHubMixin, torch.nn.Module):
 
         for adapter_name, peft_config in self.peft_config.items():
             # save only the trainable weights
-            output_state_dict = get_peft_model_state_dict(self, adapter_name, kwargs.get("state_dict", None))
+            output_state_dict = get_peft_model_state_dict(
+                self, state_dict=kwargs.get("state_dict", None), adapter_name=adapter_name
+            )
             output_dir = os.path.join(save_directory, adapter_name) if adapter_name != "default" else save_directory
             os.makedirs(output_dir, exist_ok=True)
             torch.save(output_state_dict, os.path.join(output_dir, WEIGHTS_NAME))
@@ -346,7 +348,7 @@ class PeftModel(PushToHubMixin, torch.nn.Module):
             filename, map_location=torch.device("cuda" if torch.cuda.is_available() else "cpu")
         )
         # load the weights into the model
-        set_peft_model_state_dict(self, adapter_name, adapters_weights)
+        set_peft_model_state_dict(self, adapters_weights, adapter_name=adapter_name)
         if (
             (getattr(self, "hf_device_map", None) is not None)
             and (len(set(self.hf_device_map.values()).intersection({"cpu", "disk"})) > 0)
diff --git a/src/peft/tuners/lora.py b/src/peft/tuners/lora.py
index 417bb7a..0023717 100644
--- a/src/peft/tuners/lora.py
+++ b/src/peft/tuners/lora.py
@@ -136,7 +136,7 @@ class LoraModel(torch.nn.Module):
         self.model = model
         self.forward = self.model.forward
         self.config = config
-        self.add_adapter(adapter_name)
+        self.add_adapter(adapter_name, self.config[adapter_name])
 
     def add_adapter(self, adapter_name, config=None):
         if config is not None:
diff --git a/src/peft/utils/save_and_load.py b/src/peft/utils/save_and_load.py
index fb4b252..7ebdfaf 100644
--- a/src/peft/utils/save_and_load.py
+++ b/src/peft/utils/save_and_load.py
@@ -16,7 +16,7 @@
 from .config import PeftType, PromptLearningConfig
 
 
-def get_peft_model_state_dict(model, adapter_name, state_dict=None):
+def get_peft_model_state_dict(model, state_dict=None, adapter_name="default"):
     """
     Get the state dict of the Peft model.
 
@@ -68,7 +68,7 @@ def get_peft_model_state_dict(model, adapter_name, state_dict=None):
     return to_return
 
 
-def set_peft_model_state_dict(model, adapter_name, peft_model_state_dict):
+def set_peft_model_state_dict(model, peft_model_state_dict, adapter_name="default"):
     """
     Set the state dict of the Peft model.