code: def masked_mean(vector: torch.Tensor,
                  mask: torch.Tensor,
                  dim: int,
                  keepdim: bool = False,
                  eps: float = 1e-8) -> torch.Tensor:
      
      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)
docs: 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.