Rework MPS randn fix, add randn_like fix

torch.manual_seed() already sets a CPU generator, so there is no reason to create a CPU generator manually. torch.randn_like also needs a MPS fix for k-diffusion, but a torch hijack with randn_like already exists so it can also be used for that.
This commit is contained in:
brkirch
2022-11-30 08:02:39 -05:00
parent 4d5f1691dd
commit 0fddb4a1c0
2 changed files with 8 additions and 15 deletions

View File

@@ -365,7 +365,10 @@ class TorchHijack:
if noise.shape == x.shape:
return noise
return torch.randn_like(x)
if x.device.type == 'mps':
return torch.randn_like(x, device=devices.cpu).to(x.device)
else:
return torch.randn_like(x)
# MPS fix for randn in torchsde
@@ -429,8 +432,7 @@ class KDiffusionSampler:
self.model_wrap.step = 0
self.eta = p.eta or opts.eta_ancestral
if self.sampler_noises is not None:
k_diffusion.sampling.torch = TorchHijack(self.sampler_noises)
k_diffusion.sampling.torch = TorchHijack(self.sampler_noises if self.sampler_noises is not None else [])
extra_params_kwargs = {}
for param_name in self.extra_params: