Fix proj linear count

neuralmagic · Jul 18, 2024 · 93c0d54 · 93c0d54
1 parent 415c0b7
commit 93c0d54
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/tests/test_auto_fp8.py b/tests/test_auto_fp8.py
@@ -9,7 +9,7 @@
 
 MODELS = [
     ("facebook/opt-125m", 160),
-    ("Qwen/Qwen2-0.5B-Instruct", 600),
+    ("Qwen/Qwen2-0.5B-Instruct", 620),
 ]
 
 @pytest.mark.parametrize("model_id,target_size", MODELS)
@@ -83,7 +83,7 @@ def test_kv_cache_static_quantization(model_id, target_size):
     proj_linear_count = 0
     output_scale_count = 0
     for name, _ in tensors.items():
-        if name.endswith("k_proj") or name.endswith("v_proj"):
+        if name.endswith("k_proj.weight") or name.endswith("v_proj.weight"):
             proj_linear_count += 1
         if name.endswith("k_proj.output_scale") or name.endswith("v_proj.output_scale"):
             output_scale_count += 1