# nn.TransformerEncoderLayer 3D Mask Doesn't Match the Broadcast Shape

**URL:** <https://discuss.pytorch.org/t/nn-transformerencoderlayer-3d-mask-doesnt-match-the-broadcast-shape/78880>\
**Category:** nlp\
**Created:** [April 28, 2020, 11:29pm UTC](https://discuss.pytorch.org/t/nn-transformerencoderlayer-3d-mask-doesnt-match-the-broadcast-shape/78880 "2020-04-28T23:29:02Z")\
**Posts on this page:** 3\
**Page:** 1

<div class="post-metadata">

**Author:** ![HHousen](https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/h/43a26b/32.png) [@HHousen](https://discuss.pytorch.org/u/HHousen)\
**Post date:** [April 28, 2020, 11:29pm UTC](https://discuss.pytorch.org/t/nn-transformerencoderlayer-3d-mask-doesnt-match-the-broadcast-shape/78880/1 "2020-04-28T23:29:02Z")

</div>

Hello. I am trying to use a 3D mask with a `nn.TransformerEncoderLayer`. However, I can not seem to get the dimensions right. The [`nn.Transformer` documentation](https://pytorch.org/docs/stable/nn.html#torch.nn.Transformer.forward) dictates the source size as (source sequence length, batch size, feature number) which I use below. The [`MultiHeadAttention` code](https://pytorch.org/docs/stable/_modules/torch/nn/modules/activation.html#MultiheadAttention) shows the shape of a 3D mask should be "`(N*num_heads, L, S)` where N is the batch size, L is the target sequence length, S is the source sequence length," which I also use below.

Reproducible Example Code:

```python
import torch
from torch import nn

num_heads = 8
bs = 16
num_features = 768
src_and_target_size = 21

encoder_layer = nn.TransformerEncoderLayer(num_features, num_heads, dim_feedforward=2048, dropout=0.1)

src = torch.rand(src_and_target_size, bs, num_features)

mask = torch.ones(bs*num_heads, src_and_target_size, src_and_target_size)

print(src.shape)
print(mask.shape)

encoder_layer(src, mask)

```

However, I recieve the below error message. Can you please help me find the error in my mask dimension sizes?

Error Message:

```auto
Traceback (most recent call last):
  File "test.py", line 19, in <module>
    encoder_layer(src, mask)
  File "/home/user/anaconda3/envs/test_env/lib/python3.7/site-packages/torch/nn/modules/module.py", line 532, in __call__
    result = self.forward(*input, **kwargs)
  File "/home/user/anaconda3/envs/test_env/lib/python3.7/site-packages/torch/nn/modules/transformer.py", line 283, in forward
    key_padding_mask=src_key_padding_mask)[0]
  File "/home/user/anaconda3/envs/test_env/lib/python3.7/site-packages/torch/nn/modules/module.py", line 532, in __call__
    result = self.forward(*input, **kwargs)
  File "/home/user/anaconda3/envs/test_env/lib/python3.7/site-packages/torch/nn/modules/activation.py", line 819, in forward
    attn_mask=attn_mask)
  File "/home/user/anaconda3/envs/test_env/lib/python3.7/site-packages/torch/nn/functional.py", line 3362, in multi_head_attention_forward
    attn_output_weights += attn_mask
RuntimeError: output with shape [128, 21, 21] doesn't match the broadcast shape [1, 128, 21, 21]

```

Thanks.

---

<div class="post-metadata">

**Author:** ![HHousen](https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/h/43a26b/32.png) [@HHousen](https://discuss.pytorch.org/u/HHousen)\
**Post date:** [April 28, 2020, 11:58pm UTC](https://discuss.pytorch.org/t/nn-transformerencoderlayer-3d-mask-doesnt-match-the-broadcast-shape/78880/2 "2020-04-28T23:58:02Z")

</div>

The PyTorch code that produces the error:

`torch/nn/functional.py", line 3362, in multi_head_attention_forward`:

```python
if attn_mask is not None:
    attn_mask = attn_mask.unsqueeze(0)
    attn_output_weights += attn_mask

```

Commenting out the `attn_mask.unsqueeze(0)` line makes the code run without errors. Is this a bug in PyTorch? Why add a dimension to the beginning of the attention mask?

---

<div class="post-metadata">

**Author:** ![HHousen](https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/h/43a26b/32.png) [@HHousen](https://discuss.pytorch.org/u/HHousen)\
**Post date:** [April 29, 2020, 12:11am UTC](https://discuss.pytorch.org/t/nn-transformerencoderlayer-3d-mask-doesnt-match-the-broadcast-shape/78880/3 "2020-04-29T00:11:46Z")

</div>

Solution: Upgrade to PyTorch 1.5