class_docstr: Megatron Retrieval enhanced language model
class_docstr_tok: ["Megatron","Retrieval","enhanced","language","model"]
docstr: Used for generate method only.
code: def get_forward_output_only_func(self):
      """
      Used for generate method only.
      """
  
      def fwd_output_only_func(dataloader_iter, model):
          batch = next(dataloader_iter)
          extra_arg = {}
          (
              tokens,
              attention_mask,
              retrieved,
              retrieved_mask,
              set_inference_key_value_memory,
              inference_max_sequence_len,
              neighbors,
              position_ids,
          ) = batch
  
          if len(retrieved.shape) == 1:
              retrieved = None
              retrieved_mask = None
          else:
              retrieved = retrieved.cuda()
              retrieved_mask = retrieved_mask.cuda()
  
          extra_arg[
              "set_inference_key_value_memory"
          ] = set_inference_key_value_memory[0].item()
          extra_arg[
              "inference_max_sequence_len"
          ] = inference_max_sequence_len[0].item()
          extra_arg["neighbors"] = neighbors[0].item()
          extra_arg["position_ids"] = position_ids
  
          output_tensor = model(
              tokens,
              attention_mask,
              retrieved,
              retrieved_mask,
              **extra_arg,
          )
  
          def id_func(output_tensor):
              return output_tensor, {"logits": output_tensor}
  
          return output_tensor, id_func
  
      return fwd_output_only_func
  
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 16 for calling the MegatronRetrievalModel obj's underlying member method code with example usage: obj.get_forward_output_only_func() and returns: fwd_output_only_func, output_tensor,  id_func, output_tensor,  dict