Update modelling_RW.py

modelling_RW.py

@@ -76,8 +76,8 @@ class RotaryEmbedding(torch.nn.Module):
             freqs = torch.einsum("i,j->ij", t, self.inv_freq)
             emb = torch.cat((freqs, freqs), dim=-1).to(device)
 
-            if dtype in [torch.float16, torch.bfloat16]:
-                emb = emb.float()
+            if dtype != emb.dtype:
+                emb = emb.to(dtype)
 
             self.cos_cached = emb.cos()[None, :, :]
             self.sin_cached = emb.sin()[None, :, :]
@@ -89,7 +89,7 @@ class RotaryEmbedding(torch.nn.Module):
 
     def forward(self, q, k):
         batch, seq_len, head_dim = q.shape
-        cos, sin = self.cos_sin(seq_len, q.device)
+        cos, sin = self.cos_sin(seq_len, q.device, q.dtype)
         return (q * cos) + (rotate_half(q) * sin), (k * cos) + (rotate_half(k) * sin)