diff options
author | Kohaku-Blueleaf <59680068+KohakuBlueleaf@users.noreply.github.com> | 2024-01-29 14:49:45 +0000 |
---|---|---|
committer | Kohaku-Blueleaf <59680068+KohakuBlueleaf@users.noreply.github.com> | 2024-01-29 14:49:45 +0000 |
commit | d243e24f539d717b221992e894a5db5a321bf3cd (patch) | |
tree | b4372967f300b8d209667ddc7d388421d3f274d1 /modules/devices.py | |
parent | 6e7f0860f7ae4a0ce59f9416fb9b2f3bcab44f1d (diff) | |
download | stable-diffusion-webui-gfx803-d243e24f539d717b221992e894a5db5a321bf3cd.tar.gz stable-diffusion-webui-gfx803-d243e24f539d717b221992e894a5db5a321bf3cd.tar.bz2 stable-diffusion-webui-gfx803-d243e24f539d717b221992e894a5db5a321bf3cd.zip |
Try to reverse the dtype checking mechanism
Diffstat (limited to 'modules/devices.py')
-rw-r--r-- | modules/devices.py | 7 |
1 files changed, 2 insertions, 5 deletions
diff --git a/modules/devices.py b/modules/devices.py index 8f49f7a4..f9648e9a 100644 --- a/modules/devices.py +++ b/modules/devices.py @@ -4,6 +4,7 @@ from functools import lru_cache import torch from modules import errors, shared +from modules import torch_utils if sys.platform == "darwin": from modules import mac_specific @@ -140,11 +141,7 @@ def manual_cast_forward(target_dtype): args = [arg.to(target_dtype) if isinstance(arg, torch.Tensor) else arg for arg in args] kwargs = {k: v.to(target_dtype) if isinstance(v, torch.Tensor) else v for k, v in kwargs.items()} - org_dtype = target_dtype - for param in self.parameters(): - if param.dtype != target_dtype: - org_dtype = param.dtype - break + org_dtype = torch_utils.get_param(self).dtype if org_dtype != target_dtype: self.to(target_dtype) |