code: def launch_spec_augment_kernel(
      x: torch.Tensor,
      x_len: torch.Tensor,
      freq_starts: torch.Tensor,
      freq_lengths: torch.Tensor,
      time_starts: torch.Tensor,
      time_lengths: torch.Tensor,
      freq_masks: int,
      time_masks: int,
      mask_value: float,
  ):
      """
      Helper method to launch the SpecAugment kernel
  
      Args:
          x: Pytorch tensor of shape [B, F, T] with the acoustic features.
          x_len: Pytorch tensor of shape [B] with the lengths of the padded sequence.
          freq_starts: Pytorch tensor of shape [B, M_f] with the start indices of freq masks.
          freq_widths: Pytorch tensor of shape [B, M_f] with the width of freq masks.
          time_starts: Pytorch tensor of shape [B, M_t] with the start indices of time masks.
          time_widths: Pytorch tensor of shape [B, M_t] with the width of time masks.
          freq_masks: Int value that determines the number of time masks.
          time_masks: Int value that determines the number of freq masks.
          mask_value: Float value that will be used as mask value.
  
      Returns:
          The spec augmented tensor 'x'
      """
      # Setup CUDA stream
      sh = x.shape
      stream = cuda.external_stream(
          torch.cuda.current_stream(x.device).cuda_stream
      )
  
      if time_masks > 0 or freq_masks > 0:
          # Parallelize over freq and time axis, parallel threads over batch
          # Sequential over masks (adaptive in time).
          blocks_per_grid = tuple([sh[1], sh[2]])
          # threads_per_block = min(MAX_THREAD_BUFFER, max(freq_masks, time_masks))
          threads_per_block = min(MAX_THREAD_BUFFER, x.shape[0])
  
          # Numba does not support fp16, force cast to fp32 temporarily at the expense of memory
          original_dtype = x.dtype
          cast_x = False
          if x.dtype == torch.float16:
              x = x.float()
              cast_x = True
  
          # Launch CUDA kernel
          spec_augment_kernel[
              blocks_per_grid, threads_per_block, stream, 0
          ](
              x,
              x_len,
              freq_starts,
              freq_lengths,
              time_starts,
              time_lengths,
              mask_value,
          )
          torch.cuda.synchronize()
  
          # Recast back to original dtype if earlier cast was performed
          if cast_x:
              x = x.to(dtype=original_dtype)
  
      return x
  
file_path: NeMo/nemo/collections/asr/parts/numba/spec_augment/spec_aug_numba.py
imports: import torch
  import torch.nn
  from numba import cuda
  from nemo.core.classes import Typing, typecheck
  from nemo.core.neural_types import LengthsType, NeuralType, SpectrogramType
  from nemo.utils import logging
label_desc: Use image node_id 3 for calling a global function with example usage: launch_spec_augment_kernel(x, x_len, freq_starts, freq_lengths, time_starts, time_lengths, freq_masks, time_masks, mask_value) and returns: x