diff options
| author | brkirch <brkirch@users.noreply.github.com> | 2023-05-08 16:32:40 -0400 |
|---|---|---|
| committer | brkirch <brkirch@users.noreply.github.com> | 2023-05-09 01:10:13 -0400 |
| commit | de401d8ffb46515a7cb4749f564d6a23085b4a5e (patch) | |
| tree | 4f6a8484247a3bdd6464f5c22346f962f582d997 /modules | |
| parent | 9efb809f7c2f6754367cafcce02926bf954815d5 (diff) | |
Fix generation with k-diffusion/UniPC on x64 Macs
Diffstat (limited to 'modules')
| -rw-r--r-- | modules/mac_specific.py | 5 |
1 files changed, 5 insertions, 0 deletions
diff --git a/modules/mac_specific.py b/modules/mac_specific.py index 68bffec6..40ce2101 100644 --- a/modules/mac_specific.py +++ b/modules/mac_specific.py @@ -57,3 +57,8 @@ if has_mps: # MPS workaround for https://github.com/pytorch/pytorch/issues/96113 CondFunc('torch.nn.functional.layer_norm', lambda orig_func, x, normalized_shape, weight, bias, eps, **kwargs: orig_func(x.float(), normalized_shape, weight.float() if weight is not None else None, bias.float() if bias is not None else bias, eps).to(x.dtype), lambda _, input, *args, **kwargs: len(args) == 4 and input.device.type == 'mps') + + # MPS workaround for https://github.com/pytorch/pytorch/issues/92311 + if platform.processor() == 'i386': + for funcName in ['torch.argmax', 'torch.Tensor.argmax']: + CondFunc(funcName, lambda _, input, *args, **kwargs: torch.max(input.float() if input.dtype == torch.int64 else input, *args, **kwargs)[1], lambda _, input, *args, **kwargs: input.device.type == 'mps')
\ No newline at end of file |