Update modeling_neollm.py

modeling_neollm.py

@@ -588,7 +588,8 @@ class NeoLLMRotaryEmbedding(nn.Module):
             cos = emb.cos() * self.attention_scaling
             sin = emb.sin() * self.attention_scaling
     
-        return cos, sin
+        return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
+
 
 
 def rotate_half(x):