fix

intel-analytics · Dec 24, 2024 · aabf224 · aabf224
1 parent b0338c5
commit aabf224
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/python/llm/src/ipex_llm/transformers/npu_models/mp_models_base.py b/python/llm/src/ipex_llm/transformers/npu_models/mp_models_base.py
@@ -492,7 +492,7 @@ def rotate_half(self, x, *, num_heads, seq_len, head_dim):
     def apply_rotary_pos_emb(self, *, q, k, cos, sin, position_ids,
                              num_heads, seq_len, head_dim):
         if position_ids is not None:
-            position_ids = self.squeeze(position_ids)
+            position_ids = self.reshape(position_ids, [-1])
             cos = self.gather(cos, self.convert_to_int32(position_ids), self.constant(1), 0)
             sin = self.gather(sin, self.convert_to_int32(position_ids), self.constant(1), 0)
             cos = self.unsqueeze(cos, [1])