open-mmlab/mmdetection · error · RuntimeError

The size of the 3D attn_mask is not correct.

Error message

The size of the 3D attn_mask is not correct.

What it means

A 3D attn_mask passed to MultiheadAttention.forward_attn must have shape exactly [bs*num_heads, len_q, len_k] because it is applied per-head after reshaping. Any other shape raises RuntimeError.

Source

Thrown at mmdet/models/layers/transformer/utils.py:578

                   'Only float, byte, and bool types are supported for \
                    attn_mask'

            if attn_mask.dtype == torch.uint8:
                warnings.warn('Byte tensor for attn_mask is deprecated.\
                     Use bool tensor instead.')
                attn_mask = attn_mask.to(torch.bool)
            if attn_mask.dim() == 2:
                attn_mask = attn_mask.unsqueeze(0)
                if list(attn_mask.size()) != [1, query.size(1), key.size(1)]:
                    raise RuntimeError(
                        'The size of the 2D attn_mask is not correct.')
            elif attn_mask.dim() == 3:
                if list(attn_mask.size()) != [
                        bs * self.num_heads,
                        query.size(1),
                        key.size(1)
                ]:
                    raise RuntimeError(
                        'The size of the 3D attn_mask is not correct.')
            else:
                raise RuntimeError(
                    "attn_mask's dimension {} is not supported".format(
                        attn_mask.dim()))
        # attn_mask's dim is 3 now.

        if key_padding_mask is not None and key_padding_mask.dtype == int:
            key_padding_mask = key_padding_mask.to(torch.bool)

        q = q.contiguous().view(bs, tgt_len, self.num_heads,
                                head_dims).permute(0, 2, 1, 3).flatten(0, 1)
        if k is not None:
            k = k.contiguous().view(bs, src_len, self.num_heads,
                                    head_dims).permute(0, 2, 1,
                                                       3).flatten(0, 1)
        if v is not None:
            v = v.contiguous().view(bs, src_len, self.num_heads,

View on GitHub (pinned to cfd5d3a985)

Solutions

1. Repeat/expand the mask to [bs*num_heads, len_q, len_k]: attn_mask = attn_mask.repeat_interleave(num_heads, dim=0)
2. Verify module num_heads matches what the mask was built with
3. Or use a 2D mask [len_q, len_k] broadcast over batch and heads

Example fix

# before
attn_mask = torch.zeros(bs, q_len, k_len)
# after
attn_mask = torch.zeros(bs, q_len, k_len).repeat_interleave(num_heads, dim=0)

expected = (query.size(0) * num_heads, query.size(1), key.size(1))
assert attn_mask.shape == expected, (attn_mask.shape, expected)

When it happens

Trigger: Passing attn_mask of shape [bs, len_q, len_k] (forgot the head dimension) or [bs*num_heads, len_k, len_k] (used key length for rows).

Common situations: Porting masks from vanilla nn.MultiheadAttention (which accepts [bs*num_heads, L, S]); mismatched num_heads between the mask producer and this attention module.

Related errors

• The size of the 2D attn_mask is not correct.
• attn_mask's dimension {} is not supported
• Unknown pos_tensor shape(-1):{}
• Byte tensor for attn_mask is deprecated.
• The annotation file of Open Images Challenge should be a txt

AI-assisted analysis of open-mmlab/mmdetection@cfd5d3a985 (2026-08-27). Data as JSON: /api/errors/b506d86f08d0ca1a. Report an issue: GitHub.