Try to reverse the dtype checking mechanism

This commit is contained in:
Kohaku-Blueleaf 2024-01-29 22:49:45 +08:00
parent 6e7f0860f7
commit d243e24f53

View File

@ -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)