Commit c62d17ae authored by AUTOMATIC's avatar AUTOMATIC

use the new devices.has_mps() function in register_buffer for DDIM/PLMS fix for OSX

parent 526f0aa5
......@@ -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)
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment