update code

Author: yao-fengchen

lmdeploy/pytorch/backends/dlinfer/rotary_embedding.py

@@ -40,13 +40,13 @@ def _rotary_embedding_fwd(position_ids: torch.Tensor,
 class DlinferRotaryEmbeddingImpl(RotaryEmbeddingImpl, nn.Module):
     """Base rotary embedding."""
 
-    def __init__(self, dim: int, base: int = 10000, scaling_factor: float = 1.0):
+    def __init__(self, dim: int, base: float = 10000.0, scaling_factor: float = 1.0):
         super().__init__()
         self.scaling_factor = scaling_factor
         self.dim = dim
         self.base = base
         # yapf: disable
-        inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2, dtype=torch.int64, device='cuda').float() / self.dim))
+        inv_freq = 1.0 / (self.base**(torch.arange(0, self.dim, 2, dtype=torch.float, device='cuda') / self.dim))
         # yapf: enable
         self.register_buffer('inv_freq', inv_freq, persistent=False)