From 65c1f4ed547182df635efd0d9300850110a84f02 Mon Sep 17 00:00:00 2001 From: Realcat Date: Thu, 24 Apr 2025 17:40:45 +0000 Subject: [PATCH] fix: use float32 in pos embed for MPS compatibility --- vggt/heads/utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vggt/heads/utils.py b/vggt/heads/utils.py index b6bb19b..533fc8a 100644 --- a/vggt/heads/utils.py +++ b/vggt/heads/utils.py @@ -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,)