fix

intel-analytics · Dec 24, 2024 · 95f41a5 · 95f41a5
1 parent 4c7cae0
commit 95f41a5
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/python/llm/src/ipex_llm/transformers/convert.py b/python/llm/src/ipex_llm/transformers/convert.py
@@ -1912,6 +1912,7 @@ def _optimize_post(model, lightweight_bmm=False):
 
             convert_forward(model, module.MistralModel, mistral_model_forward)
             convert_forward(model, module.MistralAttention, mistral_attention_forward)
+            convert_forward(model, module.MistralSdpaAttention, mistral_attention_forward)
             convert_forward(model, module.MistralRMSNorm, rms_norm_forward)
             convert_forward(model, module.MistralMLP, mlp_silu_forward)
     elif model.config.model_type == "gemma":