Update modelling_RW.py

avoid cache to avoid https://github.com/h2oai/h2ogpt/pull/297

modelling_RW.py

@@ -71,8 +71,23 @@ class RotaryEmbedding(torch.nn.Module):
         seq_len: int,
         device="cuda",
         dtype=torch.bfloat16,
-    ) -> torch.Tensor:
-        if seq_len != self.seq_len_cached or not self.use_cache:
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        if not self.use_cache:
+            t = torch.arange(seq_len, device=device).type_as(self.inv_freq)
+            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()
+
+            cos_cached = emb.cos()[None, :, :]
+            sin_cached = emb.sin()[None, :, :]
+
+            cos_cached = cos_cached.type(dtype)
+            sin_cached = sin_cached.type(dtype)
+
+            return cos_cached, sin_cached
+        elif seq_len != self.seq_len_cached or not self.use_cache:
             self.seq_len_cached = seq_len
             t = torch.arange(seq_len, device=device).type_as(self.inv_freq)
             freqs = torch.einsum("i,j->ij", t, self.inv_freq)
@@ -87,7 +102,7 @@ class RotaryEmbedding(torch.nn.Module):
             self.cos_cached = self.cos_cached.type(dtype)
             self.sin_cached = self.sin_cached.type(dtype)
 
-        return self.cos_cached, self.sin_cached
+            return self.cos_cached, self.sin_cached
 
     def forward(self, q, k):
         batch, seq_len, head_dim = q.shape