Unrecognized attribute: axes for operator ReduceMean

I have a pytorch model which i want to export to onnx with opset version more than 18.

But im facing issue while exporting, I am attaching minimalistic script to reproduce error.


import torch 
import numpy as np
import torch.nn as nn

class MinimalExample(nn.Module):
    def __init__(self, ):
        super().__init__()
        self.linear1 = nn.Linear(100,32)
    def forward(self, x):
        x = self.linear1(x)
        return torch.mean(x, dim=(-1,-2))

torch_model = MinimalExample()
batch_size=1
    
x = torch.randn(1, 3,100, 100, requires_grad=True)
torch_out = torch_model(x)

torch.onnx.export(torch_model,               # model being run
                  x,                         # model input (or a tuple for multiple inputs)
                  "minimalist.onnx",   # where to save the model (can be a file or file-like object)
                  export_params=True,        # store the trained parameter weights inside the model file
                  opset_version=18,          # the ONNX version to export the model to
                  do_constant_folding=False,  # whether to execute constant folding for optimization
                  input_names = ['input'],   # the model's input names
                  output_names = ['output'], # the model's output names
                  dynamic_axes={'input' : {0 : 'batch_size'},    # variable length axes
                                'output' : {0 : 'batch_size'}})

import onnx
onnx_model = onnx.load("minimalist.onnx")
onnx.checker.check_model(onnx_model)
import onnxruntime

ort_session = onnxruntime.InferenceSession("minimalist.onnx", providers=["CPUExecutionProvider",]
                                               )

def to_numpy(tensor):
    return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy()

    # compute ONNX Runtime output prediction``
ort_inputs = {ort_session.get_inputs()[0].name: to_numpy(x)}
ort_outs = ort_session.run(None, ort_inputs)

try:
    np.testing.assert_allclose(to_numpy(torch_out), ort_outs[0], rtol=1e-03, atol=1e-05)
except:
    print(print(torch_out,"\n", ort_outs[0]))
print(f"Exported model has been tested with ONNXRuntime, and the result looks good on CPU!")


The code works for opset version 17 but fails on 18. Is there any way i can replace the mean function and get everything working ?