Torch.onnx.export tensor shape gets changed

While exporting model input tensor gets changed from Tensor(1, 1, 32, 151) to Tensor: (tensor(1), tensor(1), tensor(32), tensor(151)). How can I stop this behavior?

Example code is simple:

dummy_input = torch.randn([1, 1, 32, 151], device="cuda")
onnx_reader_path = 'onnx_model.onnx'
torch.onnx.export(Model.model, dummy_input, onnx_reader_path, ...)
1 Like

I have the exact same issue with a Yolov7 model export.
It’s happening somewhere in the graph, out = torch._C._create_graph_by_tracing( function. The input is still as expected before the call, but in the first call of wrapper, the in_vars are already unflattened.

I assume this could be a Pytorch 2.0 thing, what version are you using?
Any help is appreciated.