To calculate mean along certain dimensions on masked values

    Parameters
    ----------
    vector : ``torch.Tensor``
        The vector to calculate mean.
    mask : ``torch.Tensor``
        The mask of the vector. It must be broadcastable with vector.
    dim : ``int``
        The dimension to calculate mean
    keepdim : ``bool``
        Whether to keep dimension
    eps : ``float``
        A small value to avoid zero division problem.

    Returns
    -------
    A ``torch.Tensor`` of including the mean values.

def masked_mean(vector: torch.Tensor,
                mask: torch.Tensor,
                dim: int,
                keepdim: bool = False,
                eps: float = 1e-8) -> torch.Tensor:
    """
    To calculate mean along certain dimensions on masked values

    Parameters
    ----------
    vector : ``torch.Tensor``
        The vector to calculate mean.
    mask : ``torch.Tensor``
        The mask of the vector. It must be broadcastable with vector.
    dim : ``int``
        The dimension to calculate mean
    keepdim : ``bool``
        Whether to keep dimension
    eps : ``float``
        A small value to avoid zero division problem.

    Returns
    -------
    A ``torch.Tensor`` of including the mean values.
    """
    one_minus_mask = (1.0 - mask).byte()
    replaced_vector = vector.masked_fill(one_minus_mask, 0.0)

    value_sum = torch.sum(replaced_vector, dim=dim, keepdim=keepdim)
    value_count = torch.sum(mask.float(), dim=dim, keepdim=keepdim)
    return value_sum / value_count.clamp(min=eps)