code: def process_batch(batch_data, sup_data_types_set):
      batch_dict = {}
      batch_index = 0
      for name, datatype in DATA_STR2DATA_CLASS.items():
          if (
              datatype in MAIN_DATA_TYPES
              or datatype in sup_data_types_set
          ):
              batch_dict[name] = batch_data[batch_index]
              batch_index = batch_index + 1
              if issubclass(datatype, WithLens):
                  batch_dict[name + "_lens"] = batch_data[batch_index]
                  batch_index = batch_index + 1
      return batch_dict
  
imports: from enum import Enum
  from typing import Optional, Tuple
  import librosa
  import matplotlib.pylab
  import numpy
  import torch
  from einops import rearrange
  from numba import jit, prange
  from nemo.collections.tts.torch.tts_data_types import DATA_STR2DATA_CLASS, MAIN_DATA_TYPES, WithLens
  from nemo.utils import logging
  from nemo.utils.decorators import deprecated
label_desc: Use image node_id 31 for calling a global function with example usage: process_batch(batch_data, sup_data_types_set) and returns: batch_dict