Repository: pytorch
Issue: [ONNX] 2.0 regression: dynamic shapes lost for an operator
Labels: module: onnx,triaged

### 🐛 Describe the bug

Code in : https://gist.github.com/PhilCuriosity/a19ab78dfa770c3fe495069365c5a638
The full version is on Google Cloud Drive: https://drive.google.com/file/d/1TuRH3c1p2GTnNAeDjq3kZV_DuFro7rMo/view?usp=drive_link （Code in parseq/trt）
Project source code: https://github.com/baudm/parseq

**Executing export2onnx.py in torch1.13 gives the correct result.**

> (torch1.13) root@d0811d03cfb1:/workspace/trt# python3 export_onnx_nn_part_test.py
/root/miniconda3/envs/torch1.13/lib/python3.8/site-packages/timm/models/helpers.py:7: FutureWarning: Importing from timm.models.helpers is deprecated, please import via timm.models
  warnings.warn(f"Importing from {__name__} is deprecated, please import via timm.models", FutureWarning)
Lightning automatically upgraded your loaded checkpoint from v1.9.5 to v2.0.0. To apply the upgrade to your files permanently, run `python -m pytorch_lightning.utilities.upgrade_checkpoint --file ../outputs/parseq/2024-05-27_08-56-44/checkpoints/epoch=14-step=220998-val_accuracy=91.6555-val_NED=98.4783.ckpt`
{'charset_train': './dict/hsDict802.txt', 'charset_test': './dict/hsDict802.txt', 'max_label_length': 80, 'batch_size': 72, 'lr': 0.0007, 'warmup_pct': 0.075, 'weight_decay': 0.0001, 'img_size': [64, 640], 'patch_size': [8, 16], 'embed_dim': 384, 'enc_num_heads': 6, 'enc_mlp_ratio': 4, 'enc_depth': 12, 'dec_num_heads': 12, 'dec_mlp_ratio': 4, 'dec_depth': 1, 'perm_num': 6, 'perm_forward': True, 'perm_mirrored': True, 'decode_ar': True, 'refine_iters': 0, 'dropout': 0.1, 'name': 'parseq'}
834
/root/miniconda3/envs/torch1.13/lib/python3.8/site-packages/torch/__init__.py:853: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
  assert condition, message
export ./weights/parseq_encoder_ar_iter0_bs6.onnx done!!!!!!!!!!
onnx check done
finished Simplified onnx
out-shape: torch.Size([6, 10, 832])
export ./weights/parseq_decoder_ar_iter0_bs6.onnx done!!!!!!!!!!
onnx check done
finished Simplified onnx
Exported model has been tested with ONNXRuntime, and the result looks good!

In torch2.1/2.5, the correct result is not obtained, and it seems that the operator of the original dynamic shape is frozen.

> (torch2.4) root@d0811d03cfb1:/workspace/trt# python3 export_onnx_nn_part.py 
/root/miniconda3/envs/torch2.4/lib/python3.10/site-packages/timm/models/helpers.py:7: FutureWarning: Importing from timm.models.helpers is deprecated, please import via timm.models
  warnings.warn(f"Importing from {__name__} is deprecated, please import via timm.models", FutureWarning)
Lightning automatically upgraded your loaded checkpoint from v1.9.5 to v2.4.0. To apply the upgrade to your files permanently, run `python -m pytorch_lightning.utilities.upgrade_checkpoint ../outputs/parseq/2024-05-27_08-56-44/checkpoints/epoch=14-step=220998-val_accuracy=91.6555-val_NED=98.4783.ckpt`
/root/miniconda3/envs/torch2.4/lib/python3.10/site-packages/torch/__init__.py:2041: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
  assert condition, message
export ./weights/parseq_encoder_ar_iter0_bs6.onnx done!!!!!!!!!!
onnx check done
finished Simplified onnx
out-shape: torch.Size([6, 10, 832])
export ./weights/parseq_decoder_ar_iter0_bs6.onnx done!!!!!!!!!!
onnx check done
finished Simplified onnx
2024-11-01 11:26:15.207400041 [E:onnxruntime:, sequential_executor.cc:516 ExecuteKernel] Non-zero status code returned while running Reshape node. Name:'/decoder/layers.0/self_attn/Reshape_4' Status Message: /onnxruntime_src/onnxruntime/core/providers/cpu/tensor/reshape_helper.h:45 onnxruntime::ReshapeHelper::ReshapeHelper(const onnxruntime::TensorShape&, onnxruntime::TensorShapeVector&, bool) input_shape_size == size was false. The input tensor cannot be reshaped to the requested shape. Input shape:{2,6,384}, requested shape:{10,72,32}
Traceback (most recent call last):
  File "/workspace/trt/export_onnx_nn_part.py", line 384, in <module>
    ort_outs = ort_session.run(None, ort_inputs)
  File "/root/miniconda3/envs/torch2.4/lib/python3.10/site-packages/onnxruntime/capi/onnxruntime_inference_collection.py", line 220, in run
    return self._sess.run(output_names, input_feed, run_options)
onnxruntime.capi.onnxruntime_pybind11_state.RuntimeException: [ONNXRuntimeError] : 6 : RUNTIME_EXCEPTION : Non-zero status code returned while running Reshape node. Name:'/decoder/layers.0/self_attn/Reshape_4' Status Message: /onnxruntime_src/onnxruntime/core/providers/cpu/tensor/reshape_helper.h:45 onnxruntime::ReshapeHelper::ReshapeHelper(const onnxruntime::TensorShape&, onnxruntime::TensorShapeVector&, bool) input_shape_size == size was false. The input tensor cannot be reshaped to the requested shape. Input shape:{2,6,384}, requested shape:{10,72,32}

### Versions

torch1.13

> aiohappyeyeballs==2.4.3
aiohttp==3.10.10
aiosignal==1.3.1
async-timeout==4.0.3
attrs==24.2.0
Brotli @ file:///croot/brotli-split_1714483155106/work
certifi @ file:///croot/certifi_1725551672989/work/certifi
charset-normalizer @ file:///croot/charset-normalizer_1721748349566/work
click==8.1.7
coloredlogs==15.0.1
cuda-python==12.3.0
filelock==3.16.1
flatbuffers==24.3.25
frozenlist==1.5.0
fsspec==2024.10.0
huggingface-hub==0.26.1
humanfriendly==10.0
idna @ file:///croot/idna_1714398848350/work
joblib==1.4.2
lightning-utilities==0.11.8
markdown-it-py==3.0.0
mdurl==0.1.2
mkl-fft @ file:///croot/mkl_fft_1695058164594/work
mkl-random @ file:///croot/mkl_random_1695059800811/work
mkl-service==2.4.0
mpmath==1.3.0
multidict==6.1.0
nltk==3.9.1
numpy @ file:///work/mkl/numpy_and_numpy_base_1682953417311/work
onnx==1.17.0
onnx-simplifier==0.4.36
onnxruntime==1.19.2
opencv-python==4.10.0.84
packaging==24.1
pillow @ file:///croot/pillow_1721059439630/work
polygraphy==0.49.9
propcache==0.2.0
protobuf==5.28.3
Pygments==2.18.0
PySocks @ file:///tmp/build/80754af9/pysocks_1605305779399/work
pytorch-lightning==2.0.0
PyYAML==6.0.2
regex==2024.9.11
requests @ file:///croot/requests_1721410876868/work
rich==13.9.3
safetensors==0.4.5
sympy==1.13.3
tensorrt @ file:///root/TensorRT-10.5.0.18/python/tensorrt-10.5.0-cp38-none-linux_x86_64.whl#sha256=038d9bd6997533a8d59a203354a9f31a852eed028c68e29342ceae21bfb92011
timm==1.0.11
torch==1.13.1
torchaudio==0.13.1
torchmetrics==1.5.1
torchvision==0.14.1
tqdm==4.66.5
typing_extensions @ file:///croot/typing_extensions_1715268824938/work
urllib3 @ file:///croot/urllib3_1727769808118/work
yarl==1.15.2

torch2.5

> aiohappyeyeballs==2.4.3
aiohttp==3.10.10
aiosignal==1.3.1
async-timeout==4.0.3
attrs==24.2.0
certifi==2024.8.30
charset-normalizer==3.4.0
click==8.1.7
coloredlogs==15.0.1
filelock==3.16.1
flatbuffers==24.3.25
frozenlist==1.5.0
fsspec==2024.10.0
huggingface-hub==0.26.1
humanfriendly==10.0
idna==3.10
Jinja2==3.1.4
joblib==1.4.2
lightning-utilities==0.11.8
markdown-it-py==3.0.0
MarkupSafe==3.0.2
mdurl==0.1.2
mpmath==1.3.0
multidict==6.1.0
networkx==3.4.2
nltk==3.9.1
numpy==1.24.4
nvidia-cublas-cu12==12.4.5.8
nvidia-cuda-cupti-cu12==12.4.127
nvidia-cuda-nvrtc-cu12==12.4.127
nvidia-cuda-runtime-cu12==12.4.127
nvidia-cudnn-cu12==9.1.0.70
nvidia-cufft-cu12==11.2.1.3
nvidia-curand-cu12==10.3.5.147
nvidia-cusolver-cu12==11.6.1.9
nvidia-cusparse-cu12==12.3.1.170
nvidia-nccl-cu12==2.21.5
nvidia-nvjitlink-cu12==12.4.127
nvidia-nvtx-cu12==12.4.127
onnx==1.17.0
onnx-simplifier==0.4.36
onnxruntime==1.19.2
opencv-python==4.10.0.84
packaging==24.1
pillow==11.0.0
propcache==0.2.0
protobuf==5.28.3
Pygments==2.18.0
pytorch-lightning==2.4.0
PyYAML==6.0.2
regex==2024.9.11
requests==2.32.3
rich==13.9.3
safetensors==0.4.5
sympy==1.13.1
timm==1.0.11
torch==2.5.0
torchmetrics==1.5.1
torchvision==0.20.0
tqdm==4.66.5
triton==3.1.0
typing_extensions==4.12.2
urllib3==2.2.3
yarl==1.16.0