Gradio 4 + WebUI 1.10

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,10 @@ 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"""
+    if t.device.type in ['mps', 'xpu']:
+        return torch.float32
+    return torch.float64