Try to reverse the dtype checking mechanism
This commit is contained in:
parent
6e7f0860f7
commit
d243e24f53
|
@ -4,6 +4,7 @@ from functools import lru_cache
|
||||||
|
|
||||||
import torch
|
import torch
|
||||||
from modules import errors, shared
|
from modules import errors, shared
|
||||||
|
from modules import torch_utils
|
||||||
|
|
||||||
if sys.platform == "darwin":
|
if sys.platform == "darwin":
|
||||||
from modules import mac_specific
|
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]
|
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()}
|
kwargs = {k: v.to(target_dtype) if isinstance(v, torch.Tensor) else v for k, v in kwargs.items()}
|
||||||
|
|
||||||
org_dtype = target_dtype
|
org_dtype = torch_utils.get_param(self).dtype
|
||||||
for param in self.parameters():
|
|
||||||
if param.dtype != target_dtype:
|
|
||||||
org_dtype = param.dtype
|
|
||||||
break
|
|
||||||
|
|
||||||
if org_dtype != target_dtype:
|
if org_dtype != target_dtype:
|
||||||
self.to(target_dtype)
|
self.to(target_dtype)
|
||||||
|
|
Loading…
Reference in New Issue