exportting nn.MultiheadAttention to onnx fails on pytorch==1.9.0

[RunningLeon]

Hi,

Question exportting nn.MultiheadAttention to onnx fails on pytorch==1.9.0

Solution: Upgrade pytorch >=1.10

### reproduce code 
import torch
from torch import nn

emb_size = 384
num_heads = 6
batch_first = False
q = torch.rand(1025, 1, emb_size)
model = nn.MultiheadAttention(emb_size, num_heads, batch_first=batch_first)
torch.onnx.export(model, (q, q, q), 'mha.onnx', opset_version=11, do_constant_folding=True)

error log

~/anaconda3/envs/pt1.9_v2/lib/python3.7/site-packages/torch/onnx/utils.py in _model_to_graph(model, args, verbose, input_names, output_names, operator_export_type, example_outputs, _retain_param_name, do_constant_folding, _disable_torch_constant_prop, fixed_batch_size, training, dynamic_axes)
    500     if do_constant_folding and _export_onnx_opset_version in torch.onnx.constant_folding_opset_versions:
    501         params_dict = torch._C._jit_pass_onnx_constant_fold(graph, params_dict,
--> 502                                                             _export_onnx_opset_version)
    503         torch._C._jit_pass_dce_allow_deleting_nodes_with_side_effects(graph)
    504 

RuntimeError: expected scalar type Long but found Float

env

sys.platform : linux
Python : 3.7.5 (default, Oct 25 2019, 15:51:11) [GCC 7.3.0]
CUDA available : True
numpy_random_seed : 2147483648
GPU 0 : GeForce RTX 2080
CUDA_HOME : /usr/local/cuda
NVCC : Cuda compilation tools, release 11.1, V11.1.74
GCC : gcc (Ubuntu 7.5.0-3ubuntu1~18.04) 7.5.0
PyTorch : 1.9.0+cu111
PyTorch compiling details : PyTorch built with:
  - GCC 7.3
  - C++ Version: 201402
  - Intel(R) Math Kernel Library Version 2020.0.0 Product Build 20191122 for Intel(R) 64 architecture applications
  - Intel(R) MKL-DNN v2.1.2 (Git Hash 98be7e8afa711dc9b66c8ff3504129cb82013cdb)
  - OpenMP 201511 (a.k.a. OpenMP 4.5)
  - NNPACK is enabled
  - CPU capability usage: AVX2
  - CUDA Runtime 11.1
  - NVCC architecture flags: -gencode;arch=compute_37,code=sm_37;-gencode;arch=compute_50,code=sm_50;-gencode;arch=compute_60,code=sm_60;-gencode;arch=compute_70,code=sm_70;-gencode;arch=compute_75,code=sm_75;-gencode;arch=compute_80,code=sm_80;-gencode;arch=compute_86,code=sm_86
  - CuDNN 8.0.5
  - Magma 2.5.2
  - Build settings: BLAS_INFO=mkl, BUILD_TYPE=Release, CUDA_VERSION=11.1, CUDNN_VERSION=8.0.5, CXX_COMPILER=/opt/rh/devtoolset-7/root/usr/bin/c++, CXX_FLAGS= -Wno-deprecated -fvisibility-inlines-hidden -DUSE_PTHREADPOOL -fopenmp -DNDEBUG -DUSE_KINETO -DUSE_FBGEMM -DUSE_QNNPACK -DUSE_PYTORCH_QNNPACK -DUSE_XNNPACK -DSYMBOLICATE_MOBILE_DEBUG_HANDLE -O2 -fPIC -Wno-narrowing -Wall -Wextra -Werror=return-type -Wno-missing-field-initializers -Wno-type-limits -Wno-array-bounds -Wno-unknown-pragmas -Wno-sign-compare -Wno-unused-parameter -Wno-unused-variable -Wno-unused-function -Wno-unused-result -Wno-unused-local-typedefs -Wno-strict-overflow -Wno-strict-aliasing -Wno-error=deprecated-declarations -Wno-stringop-overflow -Wno-psabi -Wno-error=pedantic -Wno-error=redundant-decls -Wno-error=old-style-cast -fdiagnostics-color=always -faligned-new -Wno-unused-but-set-variable -Wno-maybe-uninitialized -fno-math-errno -fno-trapping-math -Werror=format -Wno-stringop-overflow, LAPACK_INFO=mkl, PERF_WITH_AVX=1, PERF_WITH_AVX2=1, PERF_WITH_AVX512=1, TORCH_VERSION=1.9.0, USE_CUDA=ON, USE_CUDNN=ON, USE_EXCEPTION_PTR=1, USE_GFLAGS=OFF, USE_GLOG=OFF, USE_MKL=ON, USE_MKLDNN=ON, USE_MPI=OFF, USE_NCCL=ON, USE_NNPACK=ON, USE_OPENMP=ON, 

TorchVision : 0.10.0+cu111