def pool_masked(
    source: Tensor, mask: Tensor, num_cls: int, reduce: str = "mean"
) -> Tensor:
    """Reduce image feature maps :math:`(B, C, H, W)` or :math:`(C, H, W)` according to
    an integer index given by `mask` :math:`(B, H, W)` or :math:`(H, W)`.

    Args:
        source: Float tensor of shape :math:`(B, C, H, W)` or :math:`(C, H, W)` to be
            reduced.
        mask: Integer tensor of shape :math:`(B, H, W)` or :math:`(H, W)` containing the
            integer indices.
        num_cls: The number of classes in the possible masks.

    Returns:
        A tensor of shape :math:`(B, C, num_cls)` or :math:`(C, num_cls)`.
    """
    if source.dim() == 3:
        return _mask_reduce(source, mask, reduce, num_cls)
    elif source.dim() == 4:
        return _mask_reduce_batched(source, mask, num_cls)
    else:
        raise ValueError("source must have 3 or 4 dimensions")
Reduce image feature maps :math:`(B, C, H, W)` or :math:`(C, H, W)` according to
    an integer index given by `mask` :math:`(B, H, W)` or :math:`(H, W)`.

    Args:
        source: Float tensor of shape :math:`(B, C, H, W)` or :math:`(C, H, W)` to be
            reduced.
        mask: Integer tensor of shape :math:`(B, H, W)` or :math:`(H, W)` containing the
            integer indices.
        num_cls: The number of classes in the possible masks.

    Returns:
        A tensor of shape :math:`(B, C, num_cls)` or :math:`(C, num_cls)`.
    
lightly/models/utils.py