Commit 98ca437e authored by brkirch's avatar brkirch

Refactor and instead check if mps is being used, not availability

parent 0b5dcb3d
...@@ -182,11 +182,7 @@ def register_buffer(self, name, attr): ...@@ -182,11 +182,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:
attr = attr.to(device=devices.device, dtype=(torch.float32 if devices.device.type == 'mps' else None))
if devices.has_mps():
attr = attr.to(device="mps", dtype=torch.float32)
else:
attr = attr.to(devices.device)
setattr(self, name, attr) setattr(self, name, attr)
......
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