Fix forward pass for merged model

huggingface · Oct 17, 2023 · 5b0d813 · 5b0d813
1 parent e7bd60d
commit 5b0d813
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/optimum/onnxruntime/modeling_decoder.py b/optimum/onnxruntime/modeling_decoder.py
@@ -204,7 +204,7 @@ def forward(
         loss = None
         if self.use_cache:
             if past_key_values is not None:
-                input_ids = input_ids[:, -1:]
+                input_ids = input_ids[:, -1:] if past_key_values[0][0].shape[2] != 0 else input_ids
                 # Flatten the past_key_values (no need to flatten for models using multi-query attn)
                 if self.config.model_type not in MULTI_QUERY_ATTN_MODELS:
                     past_key_values = tuple(