about export pth to onnx

[JensenHJS]

There are some errors when exporting to onnx, looking forward to your reply .

My package versions are as follows： torch 2.2.2 torchvision 0.17.2 onnx 1.16.0 onnxruntime 1.17.3

The export script is as follows：

net = UNet(n_channels=3, n_classes=args.classes, bilinear=args.bilinear) device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') logging.info(f'Loading model {args.model}')
state_dict = torch.load(args.model, map_location=device)
mask_values = state_dict.pop('mask_values', [0, 1]) net.load_state_dict(state_dict) net.eval() dummy_input = torch.randn(1, 3, 256, 256) onnx_path = 'unet.onnx' input_names = "data" output_names = "output" torch.onnx.export(net, (dummy_input), onnx_path, input_names=input_names, output_names=output_names, verbose=False)

The error is: Traceback (most recent call last): File "/home/ubuntu/CODE/Pytorch-UNet/export_onnx.py", line 53, in torch.onnx.export(net, (dummy_input), onnx_path, input_names=input_names, output_names=output_names, verbose=False) File "/home/ubuntu/miniconda3/envs/unet/lib/python3.10/site-packages/torch/onnx/utils.py", line 516, in export _export( File "/home/ubuntu/miniconda3/envs/unet/lib/python3.10/site-packages/torch/onnx/utils.py", line 1613, in _export graph, params_dict, torch_out = _model_to_graph( File "/home/ubuntu/miniconda3/envs/unet/lib/python3.10/site-packages/torch/onnx/utils.py", line 1139, in _model_to_graph graph = _optimize_graph( File "/home/ubuntu/miniconda3/envs/unet/lib/python3.10/site-packages/torch/onnx/utils.py", line 674, in _optimize_graph _C._jit_pass_onnx_set_dynamic_input_shape(graph, dynamic_axes, input_names) TypeError: _jit_pass_onnx_set_dynamic_input_shape(): incompatible function arguments. The following argument types are supported:

1. (arg0: torch::jit::Graph, arg1: Dict[str, Dict[int, str]], arg2: List[str]) -> None