class_docstr: Megatron Retrieval enhanced language model
class_docstr_tok: ["Megatron","Retrieval","enhanced","language","model"]
code: def validation_step(self, batch, batch_idx):
      prefix = "test" if self.trainer.testing else "val"
      input_tokens_id = batch["tokens"]
      input_attn_mask = batch["tokens_mask"]
      loss_mask = batch["loss_mask"]
      retrieved_ids = batch["retrieved_ids"]
      retrieved_attn_mask = batch["retrieved_emb_mask"]
      labels = batch["labels"]
      if self.cfg.get("add_position_embedding", False):
          input_position_ids = build_position_ids(input_tokens_id)
      else:
          input_position_ids = None
      loss = self(
          input_tokens_id,
          input_attn_mask,
          retrieved_ids,
          retrieved_attn_mask,
          labels=labels,
          position_ids=input_position_ids,
      )
      loss_mask = loss_mask.float()
      lm_loss = (
          torch.sum(loss.view(-1) * loss_mask.reshape(-1))
          / loss_mask.sum()
      )
      reduced_loss = average_losses_across_data_parallel_group(
          [lm_loss]
      )
      if prefix == "val":
          self.validation_step_outputs.append(reduced_loss)
      else:
          self.test_step_outputs.apped(reduced_loss)
      return reduced_loss
  
imports: import math
  from typing import Any, List, Optional, Union
  import torch
  from omegaconf import DictConfig
  from pytorch_lightning.trainer.trainer import Trainer
  from nemo.collections.nlp.data.language_modeling.megatron.data_samplers import MegatronPretrainingRandomSampler, MegatronPretrainingSampler
  from nemo.collections.nlp.data.language_modeling.megatron.retro_dataset import build_mock_train_valid_test_datasets, build_train_valid_test_datasets
  from nemo.collections.nlp.models.language_modeling.megatron_base_model import MegatronBaseModel
  from nemo.collections.nlp.modules.common.megatron.module import Float16Module
  from nemo.collections.nlp.modules.common.megatron.mup.init import normal_
  from nemo.collections.nlp.modules.common.megatron.mup.shape import set_base_shapes
  from nemo.collections.nlp.modules.common.megatron.retrieval_token_level_encoder_decoder import MegatronRetrievalTokenLevelEncoderDecoderModule
  from nemo.collections.nlp.modules.common.megatron.utils import average_losses_across_data_parallel_group, build_position_ids, get_params_for_weight_decay_optimization
  from nemo.collections.nlp.modules.common.text_generation_strategy import model_inference_strategy_dispatcher
  from nemo.collections.nlp.modules.common.text_generation_utils import generate, get_computeprob_response, get_default_length_params, get_default_sampling_params, megatron_gpt_generate
  from nemo.collections.nlp.modules.common.tokenizer_utils import get_nmt_tokenizer
  from nemo.collections.nlp.modules.common.transformer.text_generation import LengthParam, OutputType, SamplingParam, TextGeneration
  from nemo.collections.nlp.parts.nlp_overrides import GradScaler
  from nemo.utils import AppState, logging
label_desc: Use image node_id 6 for calling the MegatronRetrievalModel obj's underlying member method code with example usage: obj.validation_step(batch, batch_idx) and returns: reduced_loss