fix: use float32 in pos embed for MPS compatibility

2025-04-24 17:41:34 +00:00
parent 350abfcdc6
commit 65c1f4ed54
1 changed files with 1 additions and 1 deletions
@@ -46,7 +46,7 @@ def make_sincos_pos_embed(embed_dim: int, pos: torch.Tensor, omega_0: float = 10
    """
    assert embed_dim % 2 == 0
    device = pos.device
-    omega = torch.arange(embed_dim // 2, dtype=torch.float if device.type == "mps" else torch.double, device=device)
+    omega = torch.arange(embed_dim // 2, dtype=torch.float32 if device.type == "mps" else torch.double, device=device)
    omega /= embed_dim / 2.0
    omega = 1.0 / omega_0**omega  # (D/2,)