fix position_ids length

huggingface · Nov 15, 2023 · 9a3b400 · 9a3b400
1 parent f187713
commit 9a3b400
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/optimum/intel/generation/modeling.py b/optimum/intel/generation/modeling.py
@@ -302,7 +302,11 @@ def forward(
             if not self.use_cache:
                 past_key_values_length = 0
             else:
-                past_key_values_length = past_key_values[0][1].shape[-2]
+                past_key_values_length = (
+                    past_key_values[0].shape[-2]
+                    if model_type.replace("-", "_") in MULTI_QUERY_ATTN_MODELS
+                    else past_key_values[0][1].shape[-2]
+                )
             position_ids = torch.arange(
                 past_key_values_length, seq_length + past_key_values_length, dtype=torch.long, device=self._device
             ).unsqueeze(0)