Run pre-commit

obvslib · Apr 11, 2024 · 0b5ef81 · 0b5ef81
1 parent 9fa08e6
commit 0b5ef81
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 5 deletions.
diff --git a/obvs/patchscope_base.py b/obvs/patchscope_base.py
@@ -227,11 +227,15 @@ def n_layers(self) -> int:
 
     @property
     def n_layers_source(self) -> int:
-        return len(getattr(getattr(self.source_model, self.MODEL_TARGET), self.LAYER_TARGET))
+        return len(
+            getattr(getattr(self.source_model, self.target_base_name), self.target_layer_name),
+        )
 
     @property
     def n_layers_target(self) -> int:
-        return len(getattr(getattr(self.target_model, self.MODEL_TARGET), self.LAYER_TARGET))
+        return len(
+            getattr(getattr(self.target_model, self.target_base_name), self.target_layer_name),
+        )
 
     def compute_precision_at_1(self, estimated_probs: torch.Tensor, true_token_index):
         """

diff --git a/tests/test_patchscopes.py b/tests/test_patchscopes.py
@@ -133,11 +133,11 @@ def test_source_forward_pass_with_attention_head_patching(patchscope):
         patchscope.source.head = 0
         patchscope.source_forward_pass()
 
-        assert patchscope._source_hidden_state.value.shape[0] == 1  # Batch size, always 1
-        assert patchscope._source_hidden_state.value.shape[1] == len(
+        assert patchscope._source_hidden_state.shape[0] == 1  # Batch size, always 1
+        assert patchscope._source_hidden_state.shape[1] == len(
             patchscope.source_tokens,
         )  # Number of tokens
         assert (
-            patchscope._source_hidden_state.value.shape[2]
+            patchscope._source_hidden_state.shape[2]
             == patchscope.source_model.transformer.h[0].attn.head_dim
         )  # Head dimension