code: def forward_step(data_iterator, model: T5Model):
      """Forward training step.
  
      Args:
          data_iterator : Input data iterator
          model (T5Model): The T5 Model
      """
  
      args = get_args()
      timers = get_timers()
  
      # Get the batch.
      timers("batch generator", log_level=2).start()
      (
          tokens_enc,
          tokens_dec,
          loss_mask,
          lm_labels,
          enc_mask,
          dec_mask,
          enc_dec_mask,
      ) = get_batch(data_iterator)
      timers("batch generator").stop()
  
      # Forward model lm_labels
      output_tensor = model(
          tokens_enc,
          tokens_dec,
          enc_mask,
          dec_mask,
          enc_dec_mask,
          lm_labels=lm_labels,
      )
  
      return output_tensor, partial(loss_func, loss_mask)
  
file_path: megatron-lm/pretrain_t5.py
imports: from functools import partial
  import torch
  from torch import Tensor
  from megatron import get_args, get_timers, print_rank_0
  from megatron.core import tensor_parallel
  from megatron.core.enums import ModelType
  from megatron.data.dataset_utils import build_train_valid_test_datasets
  from megatron.core.models.T5 import T5Model
  from megatron.training import pretrain
  from megatron.utils import average_losses_across_data_parallel_group
  from megatron.arguments import core_transformer_config_from_args
  from megatron.core.transformer.spec_utils import import_module
  from megatron.core.models.T5.t5_spec import get_t5_encoder_with_transformer_engine_block_spec, get_t5_decoder_with_transformer_engine_block_spec, get_t5_encoder_with_local_block_spec, get_t5_decoder_with_local_block_spec
label_desc: Use image node_id 4 for calling a global function with example usage: forward_step(data_iterator, model) and returns: output_tensor,  partial