def attention_pool(x: torch.Tensor,
                   pool: nn.Module,
                   in_size: tuple,
                   norm: Optional[nn.Module] = None):
    """Pooling the feature tokens.

    Args:
        x (torch.Tensor): The input tensor, should be with shape
            ``(B, num_heads, L, C)`` or ``(B, L, C)``.
        pool (nn.Module): The pooling module.
        in_size (Tuple[int]): The shape of the input feature map.
        norm (nn.Module, optional): The normalization module.
            Defaults to None.
    """
    ndim = x.ndim
    if ndim == 4:
        B, num_heads, L, C = x.shape
    elif ndim == 3:
        num_heads = 1
        B, L, C = x.shape
    else:
        raise RuntimeError(f'Unsupported input dimension {x.shape}')

    H, W = in_size
    assert L == H * W

    # (B, num_heads, H*W, C) -> (B*num_heads, C, H, W)
    x = x.reshape(B * num_heads, H, W, C).permute(0, 3, 1, 2).contiguous()
    x = pool(x)
    out_size = x.shape[-2:]

    # (B*num_heads, C, H', W') -> (B, num_heads, H'*W', C)
    x = x.reshape(B, num_heads, C, -1).transpose(2, 3)

    if norm is not None:
        x = norm(x)

    if ndim == 3:
        x = x.squeeze(1)

    return x, out_size
Pooling the feature tokens.

    Args:
        x (torch.Tensor): The input tensor, should be with shape
            ``(B, num_heads, L, C)`` or ``(B, L, C)``.
        pool (nn.Module): The pooling module.
        in_size (Tuple[int]): The shape of the input feature map.
        norm (nn.Module, optional): The normalization module.
            Defaults to None.
    
mmpretrain/models/backbones/mvit.py