Fix review issues

adapter-hub · Jun 8, 2024 · f94dbeb · f94dbeb
1 parent 6c19ea0
commit f94dbeb
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 4 deletions.
diff --git a/src/adapters/methods/adapter_layer_base.py b/src/adapters/methods/adapter_layer_base.py
@@ -159,9 +159,9 @@ def freeze_adapter(self, adapter_name: str, freeze: bool = True):
             adapter_name (str): The name of the adapter to freeze/ unfreeze.
             freeze (bool, optional): Whether to freeze the adapter. Defaults to True.
         """
-        if adapter_name in self.refts:
-            self.refts[adapter_name].train(not freeze)
-            for param in self.refts[adapter_name].parameters():
+        if adapter_name in self.adapter_modules:
+            self.adapter_modules[adapter_name].train(not freeze)
+            for param in self.adapter_modules[adapter_name].parameters():
                 param.requires_grad = not freeze
 
     def get_adapter(self, adapter_name: str) -> nn.Module:

diff --git a/tests/methods/test_reft.py b/tests/methods/test_reft.py
@@ -1,4 +1,4 @@
-from adapters import DiReftConfig, LoReftConfig
+from adapters import DiReftConfig, LoReftConfig, NoReftConfig
 from transformers.testing_utils import require_torch
 
 from .base import AdapterMethodBaseTestMixin
@@ -8,6 +8,7 @@
 class ReftTestMixin(AdapterMethodBaseTestMixin):
     reft_configs_to_test = [
         (LoReftConfig(), ["refts.{name}."]),
+        (NoReftConfig(prefix_positions=2, suffix_positions=2), ["refts.{name}."]),
         (DiReftConfig(tied_weights=True), ["refts.{name}."]),
     ]