bring in latest changes

2025-05-06 21:20:25 -07:00
parent afa0332e8c
commit cf6d058963
2 changed files with 7 additions and 4 deletions
@@ -17,10 +17,11 @@ 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