diff options
-rw-r--r-- | modules/sd_hijack.py | 3 |
1 files changed, 1 insertions, 2 deletions
diff --git a/modules/sd_hijack.py b/modules/sd_hijack.py index 75b2d22d..97979d05 100644 --- a/modules/sd_hijack.py +++ b/modules/sd_hijack.py @@ -418,8 +418,7 @@ def register_buffer(self, name, attr): if type(attr) == torch.Tensor:
if attr.device != devices.device:
- # would this not break cuda when torch adds has_mps() to main version?
- if getattr(torch, 'has_mps', False):
+ if devices.has_mps():
attr = attr.to(device="mps", dtype=torch.float32)
else:
attr = attr.to(devices.device)
|