use the new devices.has_mps() function in register_buffer for DDIM/PLMS fix for OSX
This commit is contained in:
parent
526f0aa556
commit
c62d17aee3
|
@ -418,8 +418,7 @@ def register_buffer(self, name, attr):
|
||||||
if type(attr) == torch.Tensor:
|
if type(attr) == torch.Tensor:
|
||||||
if attr.device != devices.device:
|
if attr.device != devices.device:
|
||||||
|
|
||||||
# would this not break cuda when torch adds has_mps() to main version?
|
if devices.has_mps():
|
||||||
if getattr(torch, 'has_mps', False):
|
|
||||||
attr = attr.to(device="mps", dtype=torch.float32)
|
attr = attr.to(device="mps", dtype=torch.float32)
|
||||||
else:
|
else:
|
||||||
attr = attr.to(devices.device)
|
attr = attr.to(devices.device)
|
||||||
|
|
Loading…
Reference in New Issue