update hooks in tests

Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>
vllm-project · Nov 15, 2024 · 840a41b · 840a41b
1 parent d22ad02
commit 840a41b
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/tests/llmcompressor/modifiers/calibration/test_kv_cache.py b/tests/llmcompressor/modifiers/calibration/test_kv_cache.py
@@ -54,9 +54,9 @@
 def _prep_for_calibration(module: torch.nn.Module):
     if is_attention_module(module):
         module.register_forward_pre_hook(
-            calibrate_kv_cache_input_hook(), with_kwargs=True
+            calibrate_kv_cache_input_hook, with_kwargs=True
         )
-        module.register_forward_hook(calibrate_kv_cache_output_hook())
+        module.register_forward_hook(calibrate_kv_cache_output_hook)
     module.quantization_status = QuantizationStatus.CALIBRATION
 
 

diff --git a/tests/llmcompressor/observers/test_helpers.py b/tests/llmcompressor/observers/test_helpers.py
@@ -32,7 +32,7 @@ def _prep_for_input_quant_calibration(module: torch.nn.Module):
     if not quantization_scheme:
         return
 
-    module.register_forward_pre_hook(calibrate_input_hook())
+    module.register_forward_pre_hook(calibrate_input_hook)
     module.quantization_status = QuantizationStatus.CALIBRATION