class_docstr: Megatron encoder-decoder base class
class_docstr_tok: ["Megatron","encoder-decoder","base","class"]
docstr: args_idx - maps batch into index of args (with None filling gaps)
  arg_names - corresponding names for a friendly error message
  output_name - name of output (hiddens for encode, logits for decode)
  kwargs - shared arguments (non tensors)
docstr_tok: ["args_idx","-","maps","batch","into","index","of","args","(","with","None","filling","gaps",")","arg_names","-","corresponding","names","for","a","friendly","error","message","output_name","-","name","of","output","(","hiddens","for","encode",",","logits","for","decode",")","kwargs","-","shared","arguments","(","non","tensors",")"]
code: def _get_forward_output_only_func(
      self, arg_names, output_name, **kwargs
  ):
      """
      args_idx - maps batch into index of args (with None filling gaps)
      arg_names - corresponding names for a friendly error message
      output_name - name of output (hiddens for encode, logits for decode)
      kwargs - shared arguments (non tensors)
      """
  
      def fwd_output_only_func(dataloader_iter, model):
          batch = next(dataloader_iter)
          batch = [
              x.cuda(non_blocking=True) if torch.is_tensor(x) else x
              for x in batch
          ]
  
          # map batch and shared args into forward args
          args = self._build_forward_args_from_kwargs(
              args_name=arg_names, args=batch, **kwargs
          )
          output = model(*args).contiguous()
  
          def id_func(output_tensor):
              if isinstance(output_tensor, dict):
                  # handle loss of hidden transformations ("output" is the default output)
                  output_tensor = output_tensor["output"]
  
              return output_tensor, {output_name: output_tensor}
  
          return output, id_func
  
      return fwd_output_only_func
  
imports: import copy
  import functools
  import inspect
  from typing import Any, Dict, List, Optional
  import torch
  from omegaconf import OmegaConf, open_dict
  from omegaconf.dictconfig import DictConfig
  from pytorch_lightning.accelerators import CPUAccelerator
  from pytorch_lightning.trainer.trainer import Trainer
  from nemo.collections.nlp.data.language_modeling.megatron.data_samplers import MegatronPretrainingRandomSampler, MegatronPretrainingSampler
  from nemo.collections.nlp.models.language_modeling.megatron_base_model import MegatronBaseModel
  from nemo.collections.nlp.modules.common.megatron.build_model import build_model
  from nemo.collections.nlp.modules.common.megatron.module import Float16Module
  from nemo.collections.nlp.modules.common.megatron.token_level_encoder_decoder import MegatronTokenLevelEncoderDecoderModule
  from nemo.collections.nlp.modules.common.megatron.utils import ApexGuardDefaults, average_losses_across_data_parallel_group, get_params_for_weight_decay_optimization
  from nemo.collections.nlp.modules.common.text_generation_utils import compute_beam_search_len_penalty, get_sampling_token_fn
  from nemo.collections.nlp.parts.utils_funcs import get_last_rank
  from nemo.utils import AppState, logging
label_desc: Use image node_id 21 for calling the MegatronLMEncoderDecoderModel obj's underlying member method code with example usage: obj._get_forward_output_only_func(arg_names, output_name) and returns: fwd_output_only_func, output,  id_func, output_tensor,  dict