Runtime Error raised by torch.nn.modules.activation.MultiheadAttention when bias=False, batch_first=True

[shakedbr]

🐛 Describe the bug

Hi,

When creating an object of torch.nn.modules.activation.MultiheadAttention with bias=False and batch_first=True, activating evaluation mode, and calling the forward pass you get an exception:

import torch

x = torch.rand((1, 5, 10))
model = torch.nn.modules.activation.MultiheadAttention(10, 1, bias=False, batch_first=True)
model.eval()
model(x, x, x)

Traceback (most recent call last):
  File "/Users/test.py", line 376, in <module>
    model(x,x,x)
  File "/opt/homebrew/Caskroom/miniforge/base/envs/py39/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File "/opt/homebrew/Caskroom/miniforge/base/envs/py39/lib/python3.9/site-packages/torch/nn/modules/activation.py", line 1107, in forward
    elif not all([(x.is_cuda or 'cpu' in str(x.device)) for x in tensor_args]):
  File "/opt/homebrew/Caskroom/miniforge/base/envs/py39/lib/python3.9/site-packages/torch/nn/modules/activation.py", line 1107, in <listcomp>
    elif not all([(x.is_cuda or 'cpu' in str(x.device)) for x in tensor_args]):
AttributeError: 'NoneType' object has no attribute 'is_cuda'

It seems that the following lines don't handle the case where a parameter is None.

pytorch/torch/nn/modules/activation.py

Lines 1117 to 1119 in 8cb5c55

	elif not all([(x.is_cuda or 'cpu' in str(x.device)) for x in tensor_args]):
	why_not_fast_path = "some Tensor argument is neither CUDA nor CPU"
	elif torch.is_grad_enabled() and any([x.requires_grad for x in tensor_args]):

Versions

[pip3] numpy==1.23.4
[pip3] torch==1.12.1
[pip3] torch-scatter==2.0.9
[pip3] torchaudio==0.12.1
[pip3] torchvision==0.2.2
[conda] numpy 1.23.4 py39hefdcf20_0 conda-forge
[conda] pytorch 1.12.1 py3.9_0 pytorch
[conda] torch-scatter 2.0.9 pypi_0 pypi
[conda] torchaudio 0.12.1 py39_cpu pytorch
[conda] torchvision 0.2.2 py_3 pytorch

cc @jbschlosser @bhosmer @cpuhrsch @erichan1

[cpuhrsch]

Thank you for opening the issue @shakedbr - does this issue persist with newer versions of PyTorch or nightlies?

[shakedbr]

@cpuhrsch, this also happens in version 1.13.0 and in a nightly version 1.14.0.dev20221113, but in these versions, to reproduce the bug you need to include an even number of heads e.g.:

import torch

x = torch.rand((1, 5, 10))
model = torch.nn.modules.activation.MultiheadAttention(10, num_heads=2, bias=False, batch_first=True)
model.eval()
model(x, x, x)

[cpuhrsch]

Thanks for checking @shakedbr - @mikekgfb has sent a fix and it looks like it'll be included in 1.13.1.

[malfet]

@mikekgfb can you please link the fix to the issue?

[weiwangmeta]

https://github.com/pytorch/pytorch/pull/88970/files

[weiwangmeta]

Can this be closed given the above PR? cc @mikekgfb @atalman @malfet

[mikekgfb]

Also needs #88854

[weiwangmeta]

Also needs #88854

Thank you Michael. #89855 (comment) is the cherry-pick to release/1.13

[atalman]

closing since cherry-pick is included in the release