Update modeling_internlm2.py

modeling_internlm2.py

@@ -199,8 +199,8 @@ def apply_rotary_pos_emb(q, k, cos, sin, position_ids):
     # The first two dimensions of cos and sin are always 1, so we can `squeeze` them.
     cos = cos.squeeze(1).squeeze(0)  # [seq_len, dim]
     sin = sin.squeeze(1).squeeze(0)  # [seq_len, dim]
-    cos = cos.unsqueeze(0).unsqueeze(0).expand(len(position_ids), -1, -1, -1)
-    sin = sin.unsqueeze(0).unsqueeze(0).expand(len(position_ids), -1, -1, -1)
+    cos = cos.unsqueeze(0).unsqueeze(0).expand(-1, -1, -1, -1) #(len(position_ids), -1, -1, -1)
+    sin = sin.unsqueeze(0).unsqueeze(0).expand(-1, -1, -1, -1) #(len(position_ids), -1, -1, -1)
     if q.size(2) == 1:
         q_embed = (q * cos[:, :, -1, :]) + (rotate_half(q) * sin[:, :, -1, :])
     else: