Repository: pytorch
Issue: torch.onnx.export() 1.16 with pytorch 2.2.1 produces invalid model
Labels: module: onnx,triaged

### 🐛 Describe the bug

I have a simple PyTorch model that I attempted to convert to ONNX format. The `forward()` function consists of a call to `nn.transformer.encoder()`. The conversion to ONNX using `torch.onnx.export()` completes successfully. However, when I test the model using `onnxruntime_test`, it fails except for specific input cases.

After some experiments, I discovered that the problem lies not with onnxruntime, but with the `torch.onnx.export()`. Before converting the model to TensorFlow and then exporting, I also tried an older version of `torch.onnx.export()`. Surprisingly, it worked, and the model performed correctly also with the newer onnxruntime version.

Here is a minimal example where I create the model, run it with two different tensors, export it to ONNX, and attempt to replicate the process using ONNX.

The issue seems to be related to the dynamic axes. The 'buggy' version runs correctly for the same input dimensions as those used during export.

---
Not working: pytorch  2.2.1, onnx 1.16
Working: pytorch 1.11, onnx 1.16

And the minimal code to test:

```
#!/usr/bin/env python3

import torch.nn as nn
from torch import Tensor, rand
import torch.onnx
import onnx

onnx_model = 'MicroTest.onnx'

class Test_trans(nn.Module):
    def __init__( self, emb_size=100):
        super(Test_trans, self).__init__()
        self.transformer = nn.Transformer(emb_size, 2, 2, 2, 512, 0.1)

    def forward(self, src: Tensor):
        return self.transformer.encoder(src)

def process_one_torch(session, ten):
    print('Tensor In size:', ten.size(), end='\t')
    memory = session(ten)
    print('Mem size:', memory.size());

def process_one_onnx(session, npa):
    ortvalue = onnxruntime.OrtValue.ortvalue_from_numpy(npa)
    print('In ortvalue.shape:', ortvalue.shape(), end='\t')
    memory = session.run(None, {session.get_inputs()[0].name: ortvalue})
    print('ONNX mem.shape:', memory[0].shape)

mini = Test_trans()

c_tensor_12 = rand((12,1,100))
c_tensor_10 = rand((10,1,100))

print('################################# Torch ###################################')
process_one_torch(mini, c_tensor_12)
process_one_torch(mini, c_tensor_10)

torch.onnx.export(mini, # model being run
    c_tensor_12,        # model input (or a tuple for multiple inputs)
    onnx_model,         # 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=15,   # the ONNX version to export the model to
    do_constant_folding=False,
    input_names = ['input'],   # the model's input names
    output_names = ['output'], # the model's output names
    dynamic_axes = {'input' : {0: 'max_len'}})

print('################################# ONNX RT #################################')
import onnxruntime

session = onnxruntime.InferenceSession(onnx_model, providers=["CPUExecutionProvider"])
print('Session inputs:', session.get_inputs()[0])

process_one_onnx(session, c_tensor_12.numpy())
process_one_onnx(session, c_tensor_10.numpy())   # Crash here with newer version

```

### Versions

Versions:
Collecting environment information...
PyTorch version: 2.2.1
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: Ubuntu 22.04.3 LTS (x86_64)
GCC version: (Ubuntu 11.4.0-1ubuntu1~22.04) 11.4.0
Clang version: Could not collect
CMake version: version 3.22.1
Libc version: glibc-2.35

Python version: 3.11.9 (main, Apr 19 2024, 16:48:06) [GCC 11.2.0] (64-bit runtime)
Python platform: Linux-5.15.0-82-generic-x86_64-with-glibc2.35
Is CUDA available: False
CUDA runtime version: No CUDA
CUDA_MODULE_LOADING set to: N/A
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True
CPU:
Architecture:                       x86_64
CPU op-mode(s):                     32-bit, 64-bit
Address sizes:                      46 bits physical, 48 bits virtual
Byte Order:                         Little Endian
CPU(s):                             32
On-line CPU(s) list:                0-31
Vendor ID:                          GenuineIntel
Model name:                         Intel(R) Xeon(R) CPU E5-2650L 0 @ 1.80GHz
CPU family:                         6
Model:                              45
Thread(s) per core:                 2
Core(s) per socket:                 8
Socket(s):                          2
Stepping:                           7
CPU max MHz:                        2300.0000
CPU min MHz:                        1200.0000
BogoMIPS:                           3600.28
Flags:                              fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush dts acpi mmx fxsr sse sse2 ss ht tm pbe syscall nx pdpe1gb rdtscp lm constant_tsc arch_perfmon pebs bts rep_good nopl xtopology nonstop_tsc cpuid aperfmperf pni pclmulqdq dtes64 monitor ds_cpl vmx smx est tm2 ssse3 cx16 xtpr pdcm pcid dca sse4_1 sse4_2 x2apic popcnt tsc_deadline_timer aes xsave avx lahf_lm epb pti ssbd ibrs ibpb stibp tpr_shadow vnmi flexpriority ept vpid xsaveopt dtherm ida arat pln pts md_clear flush_l1d
Virtualization:                     VT-x
L1d cache:                          512 KiB (16 instances)
L1i cache:                          512 KiB (16 instances)
L2 cache:                           4 MiB (16 instances)
L3 cache:                           40 MiB (2 instances)
NUMA node(s):                       2
NUMA node0 CPU(s):                  0-7,16-23
NUMA node1 CPU(s):                  8-15,24-31
Vulnerability Gather data sampling: Not affected
Vulnerability Itlb multihit:        KVM: Mitigation: VMX disabled
Vulnerability L1tf:                 Mitigation; PTE Inversion; VMX conditional cache flushes, SMT vulnerable
Vulnerability Mds:                  Mitigation; Clear CPU buffers; SMT vulnerable
Vulnerability Meltdown:             Mitigation; PTI
Vulnerability Mmio stale data:      Unknown: No mitigations
Vulnerability Retbleed:             Not affected
Vulnerability Spec store bypass:    Mitigation; Speculative Store Bypass disabled via prctl and seccomp
Vulnerability Spectre v1:           Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2:           Mitigation; Retpolines, IBPB conditional, IBRS_FW, STIBP conditional, RSB filling, PBRSB-eIBRS Not affected
Vulnerability Srbds:                Not affected
Vulnerability Tsx async abort:      Not affected

Versions of relevant libraries:
[pip3] numpy==1.26.4
[pip3] onnx==1.16.0
[pip3] onnx-tool==0.9.0
[pip3] onnxruntime==1.17.1
[pip3] onnxscript==0.1.0.dev20240315
[pip3] torch==2.2.1
[pip3] torchdata==0.7.1
[pip3] torchtext==0.17.1
[conda] blas                      1.0                         mkl
[conda] mkl                       2023.1.0         h213fc3f_46344
[conda] mkl-service               2.4.0           py311h5eee18b_1
[conda] mkl_fft                   1.3.8           py311h5eee18b_0
[conda] mkl_random                1.2.4           py311hdb19cb5_0
[conda] numpy                     1.26.4          py311h08b1b3b_0
[conda] numpy-base                1.26.4          py311hf175353_0
[conda] pytorch                   2.2.1              py3.11_cpu_0    pytorch
[conda] pytorch-mutex             1.0                         cpu    pytorch
[conda] torchdata                 0.7.1                     py311    pytorch
[conda] torchtext                 0.17.1                    py311    pytorch