diff options
| author | brkirch <brkirch@users.noreply.github.com> | 2022-11-12 02:17:55 -0500 |
|---|---|---|
| committer | brkirch <brkirch@users.noreply.github.com> | 2022-11-28 21:18:51 -0500 |
| commit | 98ca437edfbf71dd956d67d37f2136b12d13be0d (patch) | |
| tree | 958588db8a9c61abb0423e363b12f59699a12451 /modules/sd_hijack.py | |
| parent | 0b5dcb3d7ce397ad38312dbfc70febe7bb42dcc3 (diff) | |
Refactor and instead check if mps is being used, not availability
Diffstat (limited to 'modules/sd_hijack.py')
| -rw-r--r-- | modules/sd_hijack.py | 6 |
1 files changed, 1 insertions, 5 deletions
diff --git a/modules/sd_hijack.py b/modules/sd_hijack.py index b824b5bf..ce583950 100644 --- a/modules/sd_hijack.py +++ b/modules/sd_hijack.py @@ -182,11 +182,7 @@ def register_buffer(self, name, attr): if type(attr) == torch.Tensor:
if attr.device != devices.device:
-
- if devices.has_mps():
- attr = attr.to(device="mps", dtype=torch.float32)
- else:
- attr = attr.to(devices.device)
+ attr = attr.to(device=devices.device, dtype=(torch.float32 if devices.device.type == 'mps' else None))
setattr(self, name, attr)
|