fix mac mps dtype must be float32

2025-05-07 12:09:39 +08:00
parent 0e05a27c2f
commit 6da343dfab
2 changed files with 8 additions and 7 deletions
@@ -17,14 +17,15 @@ def project(
    dims=[-1, -2],
 ):
    dtype = v0.dtype
-    if v0.device.type == "mps":
-        v0, v1 = v0.float(), v1.float()
-    else:
-        v0, v1 = v0.double(), v1.double()
+    device_type = v0.device.type
+    if device_type == "mps":
+        v0, v1 = v0.cpu(), v1.cpu()
+
+    v0, v1 = v0.double(), v1.double()
    v1 = torch.nn.functional.normalize(v1, dim=dims)
    v0_parallel = (v0 * v1).sum(dim=dims, keepdim=True) * v1
    v0_orthogonal = v0 - v0_parallel
-    return v0_parallel.to(dtype), v0_orthogonal.to(dtype)
+    return v0_parallel.to(dtype).to(device_type), v0_orthogonal.to(dtype).to(device_type)


 def apg_forward(