Repository: pytorch
Issue: torch.jit.script failed to compile nn.MultiheadAttention when specifying the kdim and vdim parameters.
Labels: oncall: jit

### 🐛 Describe the bug

nn.MultiheadAttention cannot be correctly compiled by torch.jit.script when kdim and vdim are specified. Minimum code to reproduce the behavior:

```
import torch


mha0 = torch.nn.MultiheadAttention(256, 4)
script_mha0 = torch.jit.script(mha0)
print("Successful compiled mha0")

mha1 = torch.nn.MultiheadAttention(256, 4, kdim=128, vdim=128)
script_mha1 = torch.jit.script(mha1)
print("Successful compiled mha1")
```

Full error stack:

```
Successful compiled mha0
Traceback (most recent call last):
  File "tests/bin/jit_mha.py", line 9, in <module>
    script_mha1 = torch.jit.script(mha1)
  File "/storage09/yuxinyuan02/anaconda3/envs/pytorch/lib/python3.6/site-packages/torch/jit/_script.py", line 1097, in script
    obj, torch.jit._recursive.infer_methods_to_compile
  File "/storage09/yuxinyuan02/anaconda3/envs/pytorch/lib/python3.6/site-packages/torch/jit/_recursive.py", line 412, in create_script_module
    return create_script_module_impl(nn_module, concrete_type, stubs_fn)
  File "/storage09/yuxinyuan02/anaconda3/envs/pytorch/lib/python3.6/site-packages/torch/jit/_recursive.py", line 478, in create_script_module_impl
    create_methods_and_properties_from_stubs(concrete_type, method_stubs, property_stubs)
  File "/storage09/yuxinyuan02/anaconda3/envs/pytorch/lib/python3.6/site-packages/torch/jit/_recursive.py", line 355, in create_methods_and_properties_from_stubs
    concrete_type._create_methods_and_properties(property_defs, property_rcbs, method_defs, method_rcbs, method_defaults)
RuntimeError: 

multi_head_attention_forward(Tensor query, Tensor key, Tensor value, int embed_dim_to_check, int num_heads, Tensor in_proj_weight, Tensor? in_proj_bias, Tensor? bias_k, Tensor? bias_v, bool add_zero_attn, float dropout_p, Tensor out_proj_weight, Tensor? out_proj_bias, bool training=True, Tensor? key_padding_mask=None, bool need_weights=True, Tensor? attn_mask=None, bool use_separate_proj_weight=False, Tensor? q_proj_weight=None, Tensor? k_proj_weight=None, Tensor? v_proj_weight=None, Tensor? static_k=None, Tensor? static_v=None) -> ((Tensor, Tensor?)):
Expected a value of type 'Tensor' for argument 'in_proj_weight' but instead found type 'NoneType'.
:
  File "/storage09/yuxinyuan02/anaconda3/envs/pytorch/lib/python3.6/site-packages/torch/nn/modules/activation.py", line 1020
    
        if not self._qkv_same_embed_dim:
            attn_output, attn_output_weights = F.multi_head_attention_forward(
                                               ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
                query, key, value, self.embed_dim, self.num_heads,
                self.in_proj_weight, self.in_proj_bias,
```

### Versions

PyTorch version: 1.9.0
Is debug build: False
CUDA used to build PyTorch: 10.2
ROCM used to build PyTorch: N/A

OS: Ubuntu 18.04.5 LTS (x86_64)
GCC version: (Ubuntu 7.5.0-3ubuntu1~18.04) 7.5.0
Clang version: Could not collect
CMake version: Could not collect
Libc version: glibc-2.17

Python version: 3.6.13 |Anaconda, Inc.| (default, Jun  4 2021, 14:25:59)  [GCC 7.5.0] (64-bit runtime)
Python platform: Linux-4.15.0-142-generic-x86_64-with-debian-buster-sid
Is CUDA available: True
CUDA runtime version: Could not collect
GPU models and configuration: GPU 0: GeForce GTX 1080 Ti
Nvidia driver version: 460.73.01
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.8.0.4
/usr/lib/x86_64-linux-gnu/libcudnn_adv_infer.so.8.0.4
/usr/lib/x86_64-linux-gnu/libcudnn_adv_train.so.8.0.4
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_infer.so.8.0.4
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_train.so.8.0.4
/usr/lib/x86_64-linux-gnu/libcudnn_ops_infer.so.8.0.4
/usr/lib/x86_64-linux-gnu/libcudnn_ops_train.so.8.0.4
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] mypy-extensions==0.4.3
[pip3] numpy==1.19.2
[pip3] torch==1.9.0
[pip3] torchaudio==0.9.0a0+33b2469
[pip3] torchtext==0.10.0
[pip3] torchvision==0.10.0
[conda] blas                      1.0                         mkl  
[conda] cudatoolkit               10.2.89              hfd86e86_1  
[conda] ffmpeg                    4.3                  hf484d3e_0    pytorch
[conda] mkl                       2020.2                      256  
[conda] mkl-service               2.3.0            py36he8ac12f_0  
[conda] mkl_fft                   1.3.0            py36h54f3939_0  
[conda] mkl_random                1.1.1            py36h0573a6f_0  
[conda] mypy-extensions           0.4.3                    pypi_0    pypi
[conda] numpy                     1.19.2           py36h54aff64_0  
[conda] numpy-base                1.19.2           py36hfa32c7d_0  
[conda] pytorch                   1.9.0           py3.6_cuda10.2_cudnn7.6.5_0    pytorch
[conda] torchaudio                0.9.0                      py36    pytorch
[conda] torchtext                 0.10.0                     py36    pytorch
[conda] torchvision               0.10.0               py36_cu102    pytorch