code: def mask_along_axis(
      spec: torch.Tensor,
      spec_lengths: torch.Tensor,
      mask_width_range: Sequence[int] = (0, 30),
      dim: int = 1,
      num_mask: int = 2,
      replace_with_zero: bool = True,
  ):
      """Apply mask along the specified direction.
  
      Args:
          spec: (Batch, Length, Freq)
          spec_lengths: (Length): Not using lengths in this implementation
          mask_width_range: Select the width randomly between this range
      """
  
      org_size = spec.size()
      if spec.dim() == 4:
          # spec: (Batch, Channel, Length, Freq) -> (Batch * Channel, Length, Freq)
          spec = spec.view(-1, spec.size(2), spec.size(3))
  
      B = spec.shape[0]
      # D = Length or Freq
      D = spec.shape[dim]
      # mask_length: (B, num_mask, 1)
      mask_length = torch.randint(
          mask_width_range[0],
          mask_width_range[1],
          (B, num_mask),
          device=spec.device,
      ).unsqueeze(2)
  
      # mask_pos: (B, num_mask, 1)
      mask_pos = torch.randint(
          0,
          max(1, D - mask_length.max()),
          (B, num_mask),
          device=spec.device,
      ).unsqueeze(2)
  
      # aran: (1, 1, D)
      aran = torch.arange(D, device=spec.device)[None, None, :]
      # mask: (Batch, num_mask, D)
      mask = (mask_pos <= aran) * (aran < (mask_pos + mask_length))
      # Multiply masks: (Batch, num_mask, D) -> (Batch, D)
      mask = mask.any(dim=1)
      if dim == 1:
          # mask: (Batch, Length, 1)
          mask = mask.unsqueeze(2)
      elif dim == 2:
          # mask: (Batch, 1, Freq)
          mask = mask.unsqueeze(1)
  
      if replace_with_zero:
          value = 0.0
      else:
          value = spec.mean()
  
      if spec.requires_grad:
          spec = spec.masked_fill(mask, value)
      else:
          spec = spec.masked_fill_(mask, value)
      spec = spec.view(*org_size)
      return spec, spec_lengths
  
file_path: espnet/espnet2/layers/mask_along_axis.py
imports: import math
  from typing import Sequence, Union
  import torch
  from typeguard import check_argument_types
label_desc: Use image node_id 1 for calling a global function with example usage: mask_along_axis(spec, spec_lengths, mask_width_range, dim, num_mask, replace_with_zero) and returns: spec,  spec_lengths