Merge pull request #112 from Vincentqyw/main

Fix MPS compatibility by using float32 in positional embeddings

Author: jytime

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.float32 if device.type == "mps" else torch.double, device=device)
     omega /= embed_dim / 2.0
     omega = 1.0 / omega_0**omega  # (D/2,)