def METHOD_NAME(self, x: torch.Tensor) -> torch.Tensor:
    """Spatial pooling function.
    Args:
        x (Tensor): Input feature map.
    Returns:
        Tensor: Output tensor after spatial pooling.
    """
    batch, channel, height, width = x.size()
    if self.pooling_type == 'att':
        # [N*headers, C', H , W] C = headers * C'
        x = x.view(batch * self.n_head, self.single_header_inplanes,
                   height, width)
        input_x = x
        # [N*headers, C', H * W] C = headers * C'
        input_x = input_x.view(batch * self.n_head,
                               self.single_header_inplanes, height * width)
        # [N*headers, 1, C', H * W]
        input_x = input_x.unsqueeze(1)
        # [N*headers, 1, H, W]
        context_mask = self.conv_mask(x)
        # [N*headers, 1, H * W]
        context_mask = context_mask.view(batch * self.n_head, 1,
                                         height * width)
        # scale variance
        if self.scale_attn and self.n_head > 1:
            context_mask = context_mask / \
                           torch.sqrt(self.single_header_inplanes)
        # [N*headers, 1, H * W]
        context_mask = self.softmax(context_mask)
        # [N*headers, 1, H * W, 1]
        context_mask = context_mask.unsqueeze(-1)
        # [N*headers, 1, C', 1] =
        # [N*headers, 1, C', H * W] * [N*headers, 1, H * W, 1]
        context = torch.matmul(input_x, context_mask)
        # [N, headers * C', 1, 1]
        context = context.view(batch,
                               self.n_head * self.single_header_inplanes,
                               1, 1)
    else:
        # [N, C, 1, 1]
        context = self.avg_pool(x)
    return context