Fix style

NVIDIA · Jan 13, 2025 · d53a3cb · d53a3cb
1 parent 1d72584
commit d53a3cb
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 3 deletions.
diff --git a/evaluation/evaluate.py b/evaluation/evaluate.py
@@ -129,7 +129,7 @@ def evaluate(
     # Load press
     assert press_name in PRESS_DICT
     press = PRESS_DICT[press_name]
-    press.compression_ratio = compression_ratio
+    press.compression_ratio = compression_ratio  # type:ignore[attr-defined]
 
     # Initialize pipeline with the correct attention implementation
     model_kwargs = {"torch_dtype": "auto"}

diff --git a/kvpress/presses/key_rerotation_press.py b/kvpress/presses/key_rerotation_press.py
@@ -27,6 +27,9 @@ class KeyRerotationPress(BasePress):
 
     press: ScorerPress
 
+    def __post_init__(self):
+        assert isinstance(self.press, ScorerPress)
+
     def compress(
         self,
         module: nn.Module,
@@ -39,8 +42,6 @@ def compress(
         if self.press.compression_ratio == 0:
             return keys, values
 
-        assert isinstance(self.press, ScorerPress)
-
         # Compute scores from base press
         scores = self.press.score(module, hidden_states, keys, values, attentions, kwargs)