torch_utils.float64

return torch.float64 if device is not mps or xpu, else return torch.float32

modules/torch_utils.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import torch.nn
+import torch
 
 
 def get_param(model) -> torch.nn.Parameter:
@@ -15,3 +16,11 @@ def get_param(model) -> torch.nn.Parameter:
         return param
 
     raise ValueError(f"No parameters found in model {model!r}")
+
+
+def float64(t: torch.Tensor):
+    """return torch.float64 if device is not mps or xpu, else return torch.float32"""
+    match t.device.type:
+        case 'mps', 'xpu':
+            return torch.float32
+    return torch.float64