def aggregate_heatmaps(heatmaps: List[Tensor],
                       size: Optional[Tuple[int, int]],
                       align_corners: bool = False,
                       mode: str = 'average'):
    """Aggregate multiple heatmaps.

    Args:
        heatmaps (List[Tensor]): Multiple heatmaps to aggregate. Each should
            be in shape (B, C, H, W)
        size (Tuple[int, int], optional): The target size in (w, h). All
            heatmaps will be resized to the target size. If not given, the
            first heatmap tensor's width and height will be used as the target
            size. Defaults to ``None``
        align_corners (bool): Whether align corners when resizing heatmaps.
            Defaults to ``False``
        mode (str): Aggregation mode in one of the following:

            - ``'average'``: Get average of heatmaps. All heatmaps mush have
                the same channel number
            - ``'concat'``: Concate the heatmaps at the channel dim
    """

    if mode not in {'average', 'concat'}:
        raise ValueError(f'Invalid aggregation mode `{mode}`')

    if size is None:
        h, w = heatmaps[0].shape[2:4]
    else:
        w, h = size

    for i, _heatmaps in enumerate(heatmaps):
        assert _heatmaps.ndim == 4
        if mode == 'average':
            assert _heatmaps.shape[:2] == heatmaps[0].shape[:2]
        else:
            assert _heatmaps.shape[0] == heatmaps[0].shape[0]

        if _heatmaps.shape[2:4] != (h, w):
            heatmaps[i] = F.interpolate(
                _heatmaps,
                size=(h, w),
                mode='bilinear',
                align_corners=align_corners)

    if mode == 'average':
        output = sum(heatmaps).div(len(heatmaps))
    elif mode == 'concat':
        output = torch.cat(heatmaps, dim=1)
    else:
        raise ValueError()

    return output
Aggregate multiple heatmaps.

    Args:
        heatmaps (List[Tensor]): Multiple heatmaps to aggregate. Each should
            be in shape (B, C, H, W)
        size (Tuple[int, int], optional): The target size in (w, h). All
            heatmaps will be resized to the target size. If not given, the
            first heatmap tensor's width and height will be used as the target
            size. Defaults to ``None``
        align_corners (bool): Whether align corners when resizing heatmaps.
            Defaults to ``False``
        mode (str): Aggregation mode in one of the following:

            - ``'average'``: Get average of heatmaps. All heatmaps mush have
                the same channel number
            - ``'concat'``: Concate the heatmaps at the channel dim
    
mmpose/mmpose/models/utils/tta.py