Update chatNT.py

chatNT.py

@@ -1201,11 +1201,13 @@ class RotaryEmbeddingBis(torch.nn.Module):
     def _compute_cos_sin_tables(
         self, x: torch.Tensor, inv_freq: torch.Tensor, seq_dimension: int = 2
     ) -> tuple[torch.Tensor, torch.Tensor]:
+        print("x device : ", x.device)
         seq_len = x.shape[seq_dimension]
         # Reset the tables if the sequence length has changed,
         # or if we're on a new device (possibly due to tracing for instance)
         self._seq_len_cached = seq_len
         t = torch.arange(x.shape[seq_dimension], device=x.device).type_as(inv_freq)
+        print("t device : ", t.device)
         # freqs = torch.outer(t, inv_freq)
         freqs = torch.einsum("i, j -> ij", t, inv_freq)
 
@@ -1223,16 +1225,18 @@ class RotaryEmbeddingBis(torch.nn.Module):
     ) -> Tuple[torch.Tensor, torch.Tensor]:
         if self.rescaling_factor is None:
             inv_freq = 1.0 / (
-                self.upper_freq ** (torch.arange(0, self.dim, 2).float() / self.dim)
+                self.upper_freq ** (torch.arange(0, self.dim, 2, device=q.device).float() / self.dim)
             )
         else:
             updated_base = self.upper_freq * (
                 self.rescaling_factor ** (self.dim / (self.dim - 2))
             )
             inv_freq = 1.0 / (
-                updated_base ** (torch.arange(0, self.dim, 2).float() / self.dim)
+                updated_base ** (torch.arange(0, self.dim, 2, device=q.device).float() / self.dim)
             )
 
+        print("q device : ", q.device)
+        print("inv_freq device : ", inv_freq.device)
         self._cos_cached, self._sin_cached = self._compute_cos_sin_tables(
             q,
             inv_freq,