fix dynamic ntk device (#4483)

Author: CUHKSZzxy

lmdeploy/pytorch/backends/default/rotary_embedding.py

@@ -114,9 +114,10 @@ def __init__(self, dim: int, base: int = 10000, scaling_factor: float = 1.0, max
 
     def _ntk_inv_freq(self, seq_len: torch.Tensor):
         """ntk_inv_freq."""
+        device = seq_len.device
         base = self.base * ((self.scaling_factor * seq_len / self.max_position_embeddings) -
                             (self.scaling_factor - 1))**(self.dim / (self.dim - 2))
-        inv_freq = 1.0 / (base**(torch.arange(0, self.dim, 2, dtype=torch.int64).float() / self.dim))
+        inv_freq = 1.0 / (base**(torch.arange(0, self.dim, 2, dtype=torch.int64, device=device).float() / self.dim))
         return inv_freq
 
     def forward(self, x: torch.Tensor, position_ids: torch.Tensor):