Update modelling_RW.py

I've observed that loading this model to float16 or bfloat16 lead to bug
```
modelling_RW.py", line 289, in forward
attn_output = F.scaled_dot_product_attention(
RuntimeError: Expected query, key, and value to have the same dtype, but got query.dtype: float key.dtype: float and value.dtype: c10::Half instead.
```
This is because the output of `cos_sin` method of `RotaryEmbedding` class is `float32`.
In this commit I propose a simple fix for the model to work successfully with half precision.

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, :, :]
@@ -87,6 +87,11 @@ class RotaryEmbedding(torch.nn.Module):
 
         return self.cos_cached, self.sin_cached
 
+    def forward(self, q, k):
+        batch, seq_len, head_dim = q.shape
+        cos, sin = self.cos_sin(seq_len, q.device, q.dtype)
+        return (q * cos) + (rotate_half(q) * sin), (k * cos) + (rotate_half(k) * sin)
+
     def forward(self, q, k):
         batch, seq_len, head_dim = q.shape
         cos, sin = self.cos_sin(seq_len, q.device)