Update modeling_modern_llm.py

modeling_modern_llm.py

@@ -39,8 +39,8 @@ def rotate_half(x: torch.Tensor) -> torch.Tensor:
 
 
 def apply_rotary_pos_emb(q, k, cos, sin):
-    cos = cos.unsqueeze(0).unsqueeze(2)
-    sin = sin.unsqueeze(0).unsqueeze(2)
+    cos = cos.unsqueeze(0).unsqueeze(2).to(q.dtype)
+    sin = sin.unsqueeze(0).unsqueeze(2).to(q.dtype)
     q_embed = (q * cos) + (rotate_half(q) * sin)
     k_embed = (k * cos) + (rotate_half(k) * sin)
     return q_embed, k_embed