ONNX export: TransformerEncoder is exported with fixed input dims
Labels: module: onnx, triaged

## Issue description

When exporting to ONNX a module that contains `torch.nn.TransformerEncoder`, the time dimention is fixed to the one used during export, and the model fails to deal with arbitrary  sequence lengths.

## Code example

Given the following module:

```python
from torch import nn

class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.emb = nn.Embedding(15, 96)
        enc_layer = nn.TransformerEncoderLayer (
            96,
            nhead=12,
            dim_feedforward=96,
            dropout=0.2,
            batch_first=True
        )
        self.attn_layers = nn.TransformerEncoder(
            enc_layer,
            num_layers=10,
            enable_nested_tensor=True
        )

    def forward(self, x):
        x = self.emb(x)
        return self.attn_layers(x)
```

Exported as follows:

```python
import torch

model = Model()
torch.onnx.export(
    model,
    args=(torch.randint(0, 15, (1, 20))),
    f="model.onnx",
    opset_version=16,
    export_params=True,
    do_constant_folding=True,
    input_names=["input"],
    output_names=["output"],
    dynamic_axes={
        "input": {0: "batch_size", 1: "time"},
        "output": {0: "batch_size", 1: "time"},
    }
)
```

and ran using `onnxruntime` as follows:

```python
import numpy as np
import onnxruntime

sess = onnxruntime.InferenceSession("model.onnx")
input = np.expand_dims(
    np.array([1, 2, 3], dtype=np.int64),
    0
)
output = sess.run(None, {"input": input})
```

I get the following error:

```python
Traceback (most recent call last):
  File "D:\lab\issue_with_enc.py", line 49, in <module>
    output = sess.run(None, {"input": input})
  File "D:\lab\.env\lib\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 Reshape node. Name:'/attn_layers/layers.0/self_attn/Reshape_4' Status Message: C:\a\_work\1\s\onnxruntime\core\providers\cpu\tensor\reshape_helper.h:41 onnxruntime::ReshapeHelper::ReshapeHelper gsl::narrow_cast<int64_t>(input_shape.Size()) == size was false. The input tensor cannot be reshaped to the requested shape. Input shape:{3,1,96}, requested shape:{20,12,8}
```

## System Info

PyTorch version: 2.2.0.dev20230930+cpu
Is debug build: False
CUDA used to build PyTorch: Could not collect
ROCM used to build PyTorch: N/A

OS: Microsoft Windows 11 Pro
GCC version: Could not collect
Clang version: 15.0.2
CMake version: version 3.26.3
Libc version: N/A

Python version: 3.10.10 (tags/v3.10.10:aad5f6a, Feb  7 2023, 17:20:36) [MSC v.1929 64 bit (AMD64)] (64-bit runtime)
Python platform: Windows-10-10.0.22621-SP0
Is CUDA available: False
CUDA runtime version: Could not collect
CUDA_MODULE_LOADING set to: N/A
GPU models and configuration: GPU 0: NVIDIA GeForce MX350
Nvidia driver version: 528.49
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture=9
CurrentClockSpeed=2918
DeviceID=CPU0
Family=198
L2CacheSize=5120
L2CacheSpeed=
Manufacturer=GenuineIntel
MaxClockSpeed=2918
Name=11th Gen Intel(R) Core(TM) i7-1195G7 @ 2.90GHz
ProcessorType=3
Revision=

Versions of relevant libraries:
[pip3] numpy==1.24.3
[pip3] pytorch-lightning==2.0.6
[pip3] torch==2.2.0.dev20230930+cpu
[pip3] torchmetrics==1.0.1
[conda] Could not collect