Repository: pytorch
Issue: torch.onnx.export the torch,nn.TransformerEncoderLayer, the dynamic_axes work not right
Labels: module: onnx,triaged

### 🐛 Describe the bug

# the bug description
When I use the tensor shape like this (batch_size, seq_len, embedding_size) and i put the batch_size, seq_len in dynamic_axes to generate onnx model, I cannot use the onnx model to handle data with different seq_len, as you can see, the *code* is as follows: 
```python

#version
######
#python                    Python 3.11.1
#pytorch                   2.1.1+cu121
#onnx                       1.15.0
#onnxruntime           1.16.3
######

import torch
import torch.nn as nn
import onnxruntime

class cfg:
    # data config
    data_input_dim = 12 
    # model config
    drop_rate = 0.3
    
    # train config
    device = "cpu"
    batch_size = 512
    train_num = 100
    
    valid_freq = 1
    step_log_freq = 200
    model_save_freq = 1

    # optimizer config
    learning_rate = 5e-4
    momentum = 0.8
    weight_decay = 0.0001

class TestModel(nn.Module):
    def __init__(self):
        super(TestModel, self).__init__()
        self.pre_linear = nn.Sequential(
            nn.LayerNorm(normalized_shape=cfg.data_input_dim),
            nn.Dropout(p=cfg.drop_rate),
            nn.Linear(cfg.data_input_dim, 96),
            nn.ReLU(),
        )

        self.encoder_layer = nn.TransformerEncoderLayer(d_model=96, nhead=4, dim_feedforward=128, dropout=0.3, batch_first=True)
        self.encoder = nn.TransformerEncoder(self.encoder_layer, num_layers=5)

        self.linear = nn.Sequential(
            nn.BatchNorm1d(num_features=96),
            nn.Dropout(p=cfg.drop_rate),
            nn.Linear(96, 128),
            nn.ReLU(),
            nn.BatchNorm1d(num_features=128),
            nn.Dropout(p=cfg.drop_rate),
            nn.Linear(128, 1),
            nn.Sigmoid()
        )
    def forward(self, input_data: torch.Tensor) -> torch.Tensor: 

        data = self.pre_linear(input_data)

        transformed_data = self.encoder(data) # [batch, seq, dim]
        transformed_data = transformed_data.max(dim=-2).values  # [batch, dim]
        logit = self.linear(transformed_data)
        logit = logit[:, 0]
        logit = logit.to("cpu")
        return logit

model = TestModel()

data = torch.randn(10,5,12,device="cpu")
model.eval()
print(model(data))
input_names = ["input_data"]
output_names = ["output_data"]
dynamic_axes = {
        "input_data": {0: "batch_size", 1:"seq_len"},
        "output_data": {0: "batch_size"}
}
model.eval()
torch.onnx.export(
    model,  
    (data,),    
    "1.onnx",     
    export_params=True,     
    verbose=False,            
    training=torch.onnx.TrainingMode.EVAL,  
    input_names=input_names,    
    output_names=output_names, 
    dynamic_axes= dynamic_axes,
    )
import numpy as np

session = onnxruntime.InferenceSession('1.onnx',
                                           providers=[
                                               "CPUExecutionProvider"       # CPU
                                           ])
dummy_input = np.random.randn(2,50,12).astype(np.float32)

print(session.get_inputs())
input_name = session.get_inputs()[0].name
output_name = session.get_outputs()[0].name
print("input name: {}".format(input_name))
 
outputs = session.run([output_name], {input_name: dummy_input})

```
the *error message* is like this:
```python
2024-01-11 11:21:27.9947817 [E:onnxruntime:, sequential_executor.cc:514 onnxruntime::ExecuteKernel] Non-zero status code returned while running Reshape node. Name:'/encoder/layers.0/self_attn/Reshape_4' Status Message: C:\a\_work\1\s\onnxruntime\core\providers\cpu\tensor\reshape_helper.h:45 onnxruntime::ReshapeHelper::ReshapeHelper input_shape_size == size was false. The input tensor cannot be reshaped to the requested shape. Input shape:{50,2,96}, requested shape:{5,8,24}

Traceback (most recent call last):
  File "D:\GitCodeSource\onnx_test\test_onnx.py", line 95, in <module>
    outputs = session.run([output_name], {input_name: dummy_input})
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\venv\test\Lib\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:'/encoder/layers.0/self_attn/Reshape_4' Status Message: C:\a\_work\1\s\onnxruntime\core\providers\cpu\tensor\reshape_helper.h:45 onnxruntime::ReshapeHelper::ReshapeHelper input_shape_size == size was false. The input tensor cannot be reshaped to the requested shape. Input shape:{50,2,96}, requested shape:{5,8,24}
```

### Versions

I found it maybe one bug in the `multi_head_attention_forward` function in the `torch/nn/functional.py`,
as I change the 
```python
k = k.view(k.shape[0], bsz * num_heads, head_dim).transpose(0, 1) 
```
and
```python
v = v.view(v.shape[0], bsz * num_heads, head_dim).transpose(0, 1) 
```
in the `multi_head_attention_forward`
into 
```python
k = k.view(tgt_len, bsz * num_heads, head_dim).transpose(0, 1) 
```
and
```python
v = v.view(tgt_len, bsz * num_heads, head_dim).transpose(0, 1) 
```
The bug will fix and I can get the seq_len changable.
I guess the torch.onnx.export changes the shape[0] to a static item or the v, k shape[0] may not be as same as tgt_len so the changes I made were wrong.
Thanks for your help!