code: def get_batch(data_iterator):
      """Build the batch."""
  
      # Items and their type.
      keys = [
          "text",
          "types",
          "labels",
          "is_random",
          "loss_mask",
          "padding_mask",
      ]
      datatype = torch.int64
  
      # Broadcast data.
      if data_iterator is not None:
          data = next(data_iterator)
      else:
          data = None
      data_b = broadcast_data(keys, data, datatype)
  
      # Unpack.
      tokens = data_b["text"].long()
      types = data_b["types"].long()
      sentence_order = data_b["is_random"].long()
      loss_mask = data_b["loss_mask"].float()
      lm_labels = data_b["labels"].long()
      padding_mask = data_b["padding_mask"].long()
  
      return (
          tokens,
          types,
          sentence_order,
          loss_mask,
          lm_labels,
          padding_mask,
      )
  
file_path: ColossalAI/examples/tutorial/sequence_parallel/data/bert_helper.py
imports: import torch
  from colossalai.legacy.context import ParallelMode
  from colossalai.legacy.core import global_context
label_desc: Use image node_id 3 for calling a global function with example usage: get_batch(data_iterator) and returns: tokens,  types,  sentence_order,  loss_mask,  lm_labels,  padding_mask