ONNX support for torch.max across specific dimension

I was trying to use ONNX to export a model trained in pytorch using the max operation. ie:

max_q1, _ = torch.max(l_q1, 1)

however, I am receiving this error:

TypeError: max() got an unexpected keyword argument ‘dim’ (occurred when translating max)

Is this operation not supported by ONNX? Or, am I missing something. Thanks in advance!

1 Like