ONNX export of torch.nn.Transformer still fails
Labels: module: onnx, triaged, oncall: transformer/mha

### 🐛 Describe the bug

```python
import numpy as np
import torch

model = torch.nn.Transformer(d_model=16,
                            nhead=4,
                            num_encoder_layers=1,
                            num_decoder_layers=1,
                            batch_first=False)

src = torch.tensor(np.random.randn(3,1,16).astype(np.float32))
tgt = torch.tensor(np.random.randn(5,1,16).astype(np.float32))

torch.onnx.export(model,
                  {'src':src, 'tgt':tgt},
                  'torch_nn_transformer.onnx',
                  verbose=False,
                  input_names=['src', 'tgt'],
                  opset_version=15,
                  output_names=['out'],
                  dynamic_axes={
                    'src': {0: 'src_len'},
                    'tgt': {0: 'tgt_len'},
                    'out': {0: 'tgt_len'}
                  })
```

```
C:\Program Files\Python310\lib\site-packages\torch\nn\modules\transformer.py:282: UserWarning: enable_nested_tensor is True, but self.use_nested_tensor is False because encoder_layer.self_attn.batch_first was not True(use batch_first for better inference performance)
  warnings.warn(f"enable_nested_tensor is True, but self.use_nested_tensor is False because {why_not_sparsity_fast_path}")
Traceback (most recent call last):
  File "D:\ROOT\PythonExperiments\onnx_export_fail.py", line 13, in <module>
    torch.onnx.export(model,
  File "C:\Program Files\Python310\lib\site-packages\torch\onnx\utils.py", line 516, in export
    _export(
  File "C:\Program Files\Python310\lib\site-packages\torch\onnx\utils.py", line 1613, in _export
    graph, params_dict, torch_out = _model_to_graph(
  File "C:\Program Files\Python310\lib\site-packages\torch\onnx\utils.py", line 1135, in _model_to_graph
    graph, params, torch_out, module = _create_jit_graph(model, args)
  File "C:\Program Files\Python310\lib\site-packages\torch\onnx\utils.py", line 1011, in _create_jit_graph
    graph, torch_out = _trace_and_get_graph_from_model(model, args)
  File "C:\Program Files\Python310\lib\site-packages\torch\onnx\utils.py", line 915, in _trace_and_get_graph_from_model
    trace_graph, torch_out, inputs_states = torch.jit._get_trace_graph(
  File "C:\Program Files\Python310\lib\site-packages\torch\jit\_trace.py", line 1290, in _get_trace_graph
    outs = ONNXTracedModule(
  File "C:\Program Files\Python310\lib\site-packages\torch\nn\modules\module.py", line 1519, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "C:\Program Files\Python310\lib\site-packages\torch\nn\modules\module.py", line 1528, in _call_impl
    return forward_call(*args, **kwargs)
  File "C:\Program Files\Python310\lib\site-packages\torch\jit\_trace.py", line 138, in forward
    graph, out = torch._C._create_graph_by_tracing(
  File "C:\Program Files\Python310\lib\site-packages\torch\jit\_trace.py", line 129, in wrapper
    outs.append(self.inner(*trace_inputs))
  File "C:\Program Files\Python310\lib\site-packages\torch\nn\modules\module.py", line 1519, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "C:\Program Files\Python310\lib\site-packages\torch\nn\modules\module.py", line 1528, in _call_impl
    return forward_call(*args, **kwargs)
  File "C:\Program Files\Python310\lib\site-packages\torch\nn\modules\module.py", line 1509, in _slow_forward
    result = self.forward(*input, **kwargs)
  File "C:\Program Files\Python310\lib\site-packages\torch\nn\modules\transformer.py", line 206, in forward
    output = self.decoder(tgt, memory, tgt_mask=tgt_mask, memory_mask=memory_mask,
  File "C:\Program Files\Python310\lib\site-packages\torch\nn\modules\module.py", line 1519, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "C:\Program Files\Python310\lib\site-packages\torch\nn\modules\module.py", line 1528, in _call_impl
    return forward_call(*args, **kwargs)
  File "C:\Program Files\Python310\lib\site-packages\torch\nn\modules\module.py", line 1509, in _slow_forward
    result = self.forward(*input, **kwargs)
  File "C:\Program Files\Python310\lib\site-packages\torch\nn\modules\transformer.py", line 460, in forward
    output = mod(output, memory, tgt_mask=tgt_mask,
  File "C:\Program Files\Python310\lib\site-packages\torch\nn\modules\module.py", line 1519, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "C:\Program Files\Python310\lib\site-packages\torch\nn\modules\module.py", line 1528, in _call_impl
    return forward_call(*args, **kwargs)
  File "C:\Program Files\Python310\lib\site-packages\torch\nn\modules\module.py", line 1509, in _slow_forward
    result = self.forward(*input, **kwargs)
  File "C:\Program Files\Python310\lib\site-packages\torch\nn\modules\transformer.py", line 847, in forward
    x = self.norm2(x + self._mha_block(x, memory, memory_mask, memory_key_padding_mask, memory_is_causal))
  File "C:\Program Files\Python310\lib\site-packages\torch\nn\modules\transformer.py", line 865, in _mha_block
    x = self.multihead_attn(x, mem, mem,
  File "C:\Program Files\Python310\lib\site-packages\torch\nn\modules\module.py", line 1519, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "C:\Program Files\Python310\lib\site-packages\torch\nn\modules\module.py", line 1528, in _call_impl
    return forward_call(*args, **kwargs)
  File "C:\Program Files\Python310\lib\site-packages\torch\nn\modules\module.py", line 1509, in _slow_forward
    result = self.forward(*input, **kwargs)
  File "C:\Program Files\Python310\lib\site-packages\torch\nn\modules\activation.py", line 1241, in forward
    attn_output, attn_output_weights = F.multi_head_attention_forward(
  File "C:\Program Files\Python310\lib\site-packages\torch\nn\functional.py", line 5440, in multi_head_attention_forward
    attn_output = scaled_dot_product_attention(q, k, v, attn_mask, dropout_p, is_causal)
TypeError: scaled_dot_product_attention(): argument 'is_causal' (position 6) must be bool, not Tensor
```

BTW:
1. I'm using nightly build because otherwise I have problems with torch.nn.MultiHeadAttention export (see #110070 )
2. Mb it's a good idea to add some part of this to onnx tests. What do you think? It will show any problems with export of any part of the transformer (attention/decoder/encoder/normalization). And I think you will be getting a lot of issues like this one because of transformers popularity...

### Versions

Collecting environment information...
PyTorch version: 2.2.0.dev20230928+cpu
Is debug build: False
CUDA used to build PyTorch: Could not collect
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.27.0-rc2
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: False
CUDA runtime version: 11.8.89
CUDA_MODULE_LOADING set to: N/A
GPU models and configuration: GPU 0: NVIDIA GeForce RTX 4070 Laptop GPU
Nvidia driver version: 536.23
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture=9
CurrentClockSpeed=3601
DeviceID=CPU0
Family=107
L2CacheSize=8192
L2CacheSpeed=
Manufacturer=AuthenticAMD
MaxClockSpeed=3601
Name=AMD Ryzen 7 7745HX with Radeon Graphics
ProcessorType=3
Revision=24834

Versions of relevant libraries:
[pip3] numpy==1.26.0
[pip3] torch==2.2.0.dev20230928+cpu
[pip3] torchaudio==2.2.0.dev20230928+cpu
[pip3] torchvision==0.17.0.dev20230928+cpu
[conda] Could not collect

cc @jbschlosser @bhosmer @cpuhrsch @erichan1 @drisspg