fix

intel-analytics · Dec 25, 2024 · 6c490f2 · 6c490f2
1 parent 8bab5ae
commit 6c490f2
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/python/llm/src/ipex_llm/transformers/models/mistral.py b/python/llm/src/ipex_llm/transformers/models/mistral.py
@@ -72,7 +72,7 @@ def mistral_model_forward(
     use_quantize_kv = use_quantize_kv_cache(self.layers[0].mlp.down_proj, inputs,
                                             self.config.num_attention_heads //
                                             self.config.num_key_value_heads)
-    use_compress_kv = should_use_compresskv(inputs, inputs.size(-2)) or \
+    use_compress_kv = should_use_compresskv(inputs, inputs.size(1)) or \
         isinstance(past_key_values, DynamicCompressCache)
 
     if use_cache: