fix(onnxruntime): Add "granite" to list of model types with grouped a…

…ttention Branch: OnnxGranite Signed-off-by: Gabe Goodhart <[email protected]>
huggingface · Oct 22, 2024 · 495a3a0 · 495a3a0
1 parent 3a96b22
commit 495a3a0
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/optimum/onnxruntime/modeling_decoder.py b/optimum/onnxruntime/modeling_decoder.py
@@ -340,7 +340,7 @@ def prepare_past_key_values(
             if self.model_type == "gemma":
                 num_attention_heads = self.normalized_config.num_key_value_heads
                 embed_size_per_head = self.normalized_config.head_dim
-            elif self.model_type in {"mistral", "llama", "qwen2"}:
+            elif self.model_type in {"mistral", "llama", "qwen2", "granite"}:
                 num_attention_heads = self.normalized_config.num_key_value_heads
             else:
                 num_attention_heads = self.normalized_config.num_attention_heads