Commit de401d8f authored by brkirch's avatar brkirch
Browse files

Fix generation with k-diffusion/UniPC on x64 Macs

parent 9efb809f
Loading
Loading
Loading
Loading
+5 −0
Original line number Diff line number Diff line
@@ -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