torch_utils.float64

return torch.float64 if device is not mps or xpu, else return torch.float32
AUTOMATIC1111 · May 16, 2024 · 9c8075b · 9c8075b
1 parent ddb28b3
commit 9c8075b
Showing 1 changed file with 9 additions and 0 deletions.
diff --git a/modules/torch_utils.py b/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