bug fix

gmlwns2000 · Mar 5, 2024 · 68d93b5 · 68d93b5
1 parent a5866da
commit 68d93b5
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/vllm/model_executor/layers/rotary_embedding.py b/vllm/model_executor/layers/rotary_embedding.py
@@ -306,7 +306,9 @@ def _compute_inv_freq(self, scaling_factor: float) -> torch.Tensor:
         # Get n-d rotational scaling corrected for extrapolation
         inv_freq_mask = (1 - _yarn_linear_ramp_mask(
             low, high, self.rotary_dim // 2,
-            dtype=torch.float)) * self.extrapolation_factor
+            dtype=torch.float,
+            device=pos_freqs.device
+        )) * self.extrapolation_factor
         inv_freq = inv_freq_interpolation * (
             1 - inv_freq_mask) + inv_freq_extrapolation * inv_freq_mask
         return inv_freq