Search code examples
pythonpytorchcoremlcoremltools

Conversion pytorch to coreml for element-wise maximum operation

I tried to convert a PyTorch model to coreml with the element-wise maximum operation based on coremltools.

With torch.max operation, I got

ValueError: node input.2 (max) got 2 input(s), expected [3]

With torch.maximum operation

RuntimeError: PyTorch convert function for op 'maximum' not implemented.

Any idea to solve this issue?

Solution

  • I encountered the same issue in conversion with pytorch element-wise operation to coreml model, but solved it by adding support for torch.maximum and torch.minimum for the MIL converter with the torch frontend.

    @register_torch_op
def maximum(context, node):
    inputs = _get_inputs(context, node)
    x = inputs[0]
    y = inputs[1]
    out = mb.maximum(x=x, y=y, name=node.name)
    context.add(out)