From 9c8075ba8e538f695ef25f85e6513227b58b71ce Mon Sep 17 00:00:00 2001 From: w-e-w <40751091+w-e-w@users.noreply.github.com> Date: Thu, 16 May 2024 23:16:50 +0900 Subject: [PATCH] torch_utils.float64 return torch.float64 if device is not mps or xpu, else return torch.float32 --- modules/torch_utils.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/modules/torch_utils.py b/modules/torch_utils.py index e5b52393e..a07e02853 100644 --- 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