open-mmlab/mmdetection · error · RuntimeError

The size of the 2D attn_mask is not correct.

Error message

The size of the 2D attn_mask is not correct.

What it means

In MultiheadAttention.forward_attn, a 2D attn_mask is unsqueezed to [1, len_q, len_k]; if its shape does not exactly equal [1, query.size(1), key.size(1)] a RuntimeError is raised because the mask cannot be broadcast by the underlying attention math.

Source

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

        v = value

        if attn_mask is not None:
            assert attn_mask.dtype == torch.float32 or \
                   attn_mask.dtype == torch.float64 or \
                   attn_mask.dtype == torch.float16 or \
                   attn_mask.dtype == torch.uint8 or \
                   attn_mask.dtype == torch.bool, \
                   '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)

View on GitHub (pinned to cfd5d3a985)

Solutions

  1. Build the mask with shape [num_queries, num_keys] i.e. [query.size(1), key.size(1)]
  2. Print query.shape, key.shape and attn_mask.shape before the call and align them
  3. For per-head masks supply a 3D mask of shape [bs*num_heads, len_q, len_k]

Example fix

# before
attn_mask = torch.zeros(k, k)  # used in cross-attn with q_len != k_len
# after
attn_mask = torch.zeros(q_len, k_len)
Defensive patterns

Strategy: validation

Validate before calling

assert attn_mask.shape == (query.size(1), key.size(1)), attn_mask.shape

Type guard

def valid_2d_mask(m, q, k): return m.dim() == 2 and m.shape == (q, k)

Prevention

When it happens

Trigger: Passing attn_mask of shape [len_k, len_k] or [len_q+1, len_k] where query and key have different sequence lengths (cross-attention with text/encoder memory of different length).

Common situations: Using square masks from self-attention code in cross-attention (query len != key len); padding/truncating sequences inconsistently between queries and keys; language-guided detection where text tokens count mismatches.

Related errors


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