code: def get_batch(data_iterator):
      """Build the batch."""
      data = next(data_iterator)
  
      # only data parallelism; no need for broadcast
      images = data[0][0].cuda()
      masks = data[0][1].cuda()
      return images, masks
  
file_path: megatron-lm/pretrain_vision_inpaint.py
imports: import torch
  import torch.nn.functional
  from functools import partial
  from megatron import get_args, get_timers, print_rank_0, print_rank_last
  from megatron.core.enums import ModelType
  from megatron.data.vit_dataset import build_train_valid_datasets
  from megatron.model.vision.inpainting import VitInpaintingModel
  from megatron.model.vision.inpainting import MitInpaintingModel
  from megatron.training import pretrain
  from megatron.utils import average_losses_across_data_parallel_group
  from tasks.vision.segmentation.metrics import SSIM, PSNR
  from megatron.arguments import core_transformer_config_from_args
label_desc: Use image node_id 2 for calling a global function with example usage: get_batch(data_iterator) and returns: images,  masks