code: def masked_mean(
      tensor: torch.Tensor, mask: torch.Tensor, dim: int = 1
  ) -> torch.Tensor:
      tensor = tensor * mask
      tensor = tensor.sum(dim=dim)
      mask_sum = mask.sum(dim=dim)
      mean = tensor / (mask_sum + 1e-8)
      return mean
  
file_path: ColossalAI/applications/Chat/coati/models/utils.py
imports: from typing import Optional, Union
  import torch
  import torch.nn.functional
label_desc: Use image node_id 5 for calling a global function with example usage: masked_mean(tensor, mask, dim) and returns: mean