It looks related to the max pools being instantiated in forward. I would suggest instantiating them in the __init__ and then using them in forward to see if that address the issue with the export to ONNX.