Repository: pytorch
Issue: Can't export an onnx with dynamic axes if it contains backpropagation
Labels: module: onnx,triaged

### 🐛 Describe the bug

This script will reproduce the issue I'm experiencing:

```
import torch
import onnxruntime as ort

doBackprop = True

class Model(torch.nn.Module):

    def forward(self, inputs):

        if doBackprop:
            xseq = inputs[0].detach().requires_grad_()
            with torch.enable_grad():

                y = torch.pow(xseq, 2).sum(dim=-1).mean(dim=-1)
                loss = y.sum() + y.mean()
                
                return torch.autograd.grad([loss], [xseq])[0]
        else:
            xseq = inputs[0]
            return torch.pow(xseq, 2).sum()


# export onnx:
torch.onnx.export(
    Model(),
    [torch.randn([2, 2, 10]) ** 2],
    "modelthing.onnx",
    input_names=["x"],
    dynamic_axes={'x': {0: 'batch_size'}},
    opset_version=17,
    output_names=["lossgrad"],
    verbose=True
)

# load and evaluate with a different shape:
ort_session = ort.InferenceSession("modelthing.onnx")

print("run onnx with export shape")
ort_session.run(None, { "x":torch.randn([2, 2, 10]).numpy() })
print("success")

print("run onnx with different batch size")
ort_session.run(None, { "x":torch.randn([4, 2, 10]).numpy() })
print("success")
```

Running it with doBackprop set to True will give you this error, indicating that something is erroneously hard coding the tensor shapes internally:

```
Traceback (most recent call last):
  File "C:\Git\pytorch\dynamicshapesbroken.py", line 45, in <module>
    ort_session.run(None, { "x":torch.randn([4, 2, 10]).numpy() })
  File "C:\Users\dminor\AppData\Roaming\Python\Python310\site-packages\onnxruntime\capi\onnxruntime_inference_collection.py", line 217, 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 Mul node. Name:'/Mul_2' Status Message: C:\a\_work\1\s\onnxruntime\core/providers/cpu/math/element_wise_ops.h:560 onnxruntime::BroadcastIterator::Append axis == 1 || axis == largest was false. Attempting to broadcast an axis by a dimension other than 1. 2 by 4
```

If you run this script with the doBackprop variable set to False, you should get the following output, indicating that it's loaded the exported onnx and successfully run it with different batch sizes:

```
run onnx with export shape
success
run onnx with different batch size
success
```

Not sure why the autograd stuff is breaking dynamic axes, but I need to get it fixed

### Versions

PyTorch version: 2.2.1+cu118
Is debug build: False
CUDA used to build PyTorch: 11.8
ROCM used to build PyTorch: N/A

OS: Microsoft Windows 10 Enterprise
GCC version: Could not collect
Clang version: Could not collect
CMake version: version 3.25.1
Libc version: N/A

Python version: 3.10.11 (tags/v3.10.11:7d4cc5a, Apr 5 2023, 00:38:17) [MSC v.1929 64 bit (AMD64)] (64-bit runtime)
Python platform: Windows-10-10.0.19045-SP0
Is CUDA available: True
CUDA runtime version: 10.0.130
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: GPU 0: NVIDIA GeForce RTX 3090 Ti
Nvidia driver version: 536.23
cuDNN version: C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v11.8\bin\cudnn_ops_train64_8.dll
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture=9
CurrentClockSpeed=3501
DeviceID=CPU0
Family=107
L2CacheSize=16384
L2CacheSpeed=
Manufacturer=AuthenticAMD
MaxClockSpeed=3501
Name=AMD Ryzen Threadripper PRO 3975WX 32-Cores
ProcessorType=3
Revision=12544

Versions of relevant libraries:
[pip3] functorch==2.0.0
[pip3] lovely-numpy==0.2.8
[pip3] numpy==1.24.3
[pip3] onnx==1.14.1
[pip3] onnx-graphsurgeon==0.3.27
[pip3] onnxconverter-common==1.13.0
[pip3] onnxruntime==1.15.1
[pip3] optree==0.10.0
[pip3] pytorch-lightning==1.4.2
[pip3] tf2onnx==1.16.0
[pip3] torch==2.2.1+cu118
[pip3] torch-cluster==1.6.1
[pip3] torch-fidelity==0.3.0
[pip3] torch-geometric==2.3.0
[pip3] torch-scatter==2.1.1
[pip3] torch-sparse==0.6.17
[pip3] torch-spline-conv==1.2.2
[pip3] torchaudio==2.2.1+cu118
[pip3] torchmetrics==0.6.0
[pip3] torchvision==0.17.1+cu118
[conda] blas 1.0 mkl
[conda] cudatoolkit 11.6.0 hc0ea762_10 conda-forge
[conda] libblas 3.9.0 16_win64_mkl conda-forge
[conda] libcblas 3.9.0 16_win64_mkl conda-forge
[conda] liblapack 3.9.0 16_win64_mkl conda-forge
[conda] mkl 2022.1.0 h6a75c08_874 conda-forge
[conda] mkl-include 2023.2.0 intel_49496 intel
[conda] mkl-static 2023.2.0 intel_49496 intel
[conda] numpy 1.24.2 py310hd02465a_0 conda-forge
[conda] pytorch 1.12.0 py3.10_cuda11.6_cudnn8_0 pytorch
[conda] pytorch-mutex 1.0 cuda pytorch