Update rotary.py

rotary.py

@@ -534,7 +534,7 @@ class RotaryEmbedding(torch.nn.Module):
             # And the output of arange can be quite large, so bf16 would lose a lot of precision.
             # However, for compatibility reason, we add an option to use the dtype of self.inv_freq.
             if rotary_base_changed:
-                self.inv_freq = self._compute_inv_freq(device=self.inv_freq.device)
+                self.inv_freq = self._compute_inv_freq(device=device)
             if self.pos_idx_in_fp32:
                 t = torch.arange(seqlen, device=device, dtype=torch.float32)
                 # We want fp32 here as well since inv_freq will be multiplied with t, and the output