Merge pull request #5194 from brkirch/autocast-and-mps-randn-fixes

Use devices.autocast() and fix MPS randn issues
This commit is contained in:
AUTOMATIC1111
2022-12-03 09:58:08 +03:00
committed by GitHub
8 changed files with 29 additions and 31 deletions

View File

@@ -183,11 +183,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)