diff --git a/vggt/heads/utils.py b/vggt/heads/utils.py index d7af1f6..b6bb19b 100644 --- a/vggt/heads/utils.py +++ b/vggt/heads/utils.py @@ -45,7 +45,8 @@ def make_sincos_pos_embed(embed_dim: int, pos: torch.Tensor, omega_0: float = 10 - emb: The generated 1D positional embedding. """ assert embed_dim % 2 == 0 - omega = torch.arange(embed_dim // 2, dtype=torch.double, device=pos.device) + device = pos.device + omega = torch.arange(embed_dim // 2, dtype=torch.float if device.type == "mps" else torch.double, device=device) omega /= embed_dim / 2.0 omega = 1.0 / omega_0**omega # (D/2,)