Fix bottleneck average composition computation (#590)

adapter-hub · Oct 12, 2023 · 6da07d1 · 6da07d1
1 parent 5652c0b
commit 6da07d1
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/src/adapters/layer.py b/src/adapters/layer.py
@@ -631,8 +631,8 @@ def adapter_average_output(self, adapter_setup: Average, hidden_states, input_te
                 )
             # Case X: No adapter which is part of this module -> ignore
 
-        weights = torch.tensor(adapter_setup.weights).unsqueeze(1).unsqueeze(1).to(hidden_states.device)
-        hidden_states = torch.mean(torch.cat(children_hidden, 0) * weights, 0)
+        weights = torch.tensor(adapter_setup.weights)[:, None, None, None].to(hidden_states.device)
+        hidden_states = torch.mean(torch.stack(children_hidden, 0) * weights, 0)
 
         return hidden_states
 

diff --git a/tests_adapters/composition/test_adapter_composition.py b/tests_adapters/composition/test_adapter_composition.py
@@ -224,9 +224,9 @@ def test_average(self):
         model.set_active_adapters(Average("a", "b", "c", "d"))
 
         inputs = {}
-        inputs["input_ids"] = ids_tensor((1, 128), 1000)
+        inputs["input_ids"] = ids_tensor((2, 128), 1000)
         logits = model(**inputs).logits
-        self.assertEqual(logits.shape, (1, 2))
+        self.assertEqual(logits.shape, (2, 2))
 
 
 class PrefixTuningCompositionTest(AdapterCompositionTest):