code: def forward_step(batch, model, eval_metric, config):
      """Forward step."""
  
      # Get the batch.
      (
          tokens,
          labels,
          attention_mask,
          position_ids,
          loss_mask,
      ) = process_batch(batch)
  
      # Tell the model what our actual batch size will be
      args = get_args()
      args.micro_batch_size = len(labels)
  
      tensor_shape = (
          args.seq_length,
          args.micro_batch_size,
          args.hidden_size,
      )
      input_tensor = recv_forward(tensor_shape, config)
  
      # Forward pass through the model.
      unwrapped_model = unwrap_model(model)
      unwrapped_model.set_input_tensor(input_tensor)
      output = model(tokens, position_ids, attention_mask)
  
      send_forward(output, config)
  
      if parallel_state.is_pipeline_last_stage():
          # For loss, return the unreduced loss.
          if eval_metric == "loss":
              losses = tensor_parallel.vocab_parallel_cross_entropy(
                  output.contiguous().float(), labels.contiguous()
              )
              loss = torch.sum(
                  losses.view(-1)
                  * loss_mask.contiguous().view(-1).float()
              )
              return loss
  
          # For accuracy, return the number of correctly predicted samples.
          if eval_metric == "accuracy":
              outputs = torch.argmax(output, -1)
              correct = (outputs == labels).float()
              correct[(1 - loss_mask).bool()] = 1
              correct = correct.prod(-1)
              return correct.sum()
  
          raise NotImplementedError(
              "forward method for evaluation metric {} "
              "is not implemented.".format(eval_metric)
          )
      return None
  
file_path: megatron-lm/tasks/zeroshot_gpt/evaluate.py
imports: import math
  import torch
  from megatron import get_args
  from megatron import print_rank_0, is_last_rank
  from megatron import get_tokenizer
  from megatron.core import parallel_state, tensor_parallel
  from megatron.checkpointing import load_checkpoint
  from megatron.model import GPTModel
  from megatron.training import get_model
  from megatron.utils import get_ltor_masks_and_position_ids, unwrap_model
  from megatron.core.pipeline_parallel.p2p_communication import recv_forward, send_forward
  from megatron.arguments import core_transformer_config_from_args
  from tasks.finetune_utils import build_data_loader
  from .datasets import build_dataset
label_desc: Use image node_id 3 for calling a global function with example usage: forward_step(batch, model, eval_metric, config) and returns: None, loss, correct