class_docstr: Megatron Bert pretraining.
  Model returns [batch, seq, hidden] shape
class_docstr_tok: ["Megatron","Bert","pretraining",".","Model","returns","[","batch",",","seq",",","hidden","]","shape"]
code: def get_forward_output_and_loss_func(self):
      def fwd_output_and_loss_func(
          dataloader_iter, model, checkpoint_activations_all_layers=None
      ):
          if (
              parallel_state.get_pipeline_model_parallel_world_size()
              == 1
          ):
              batch = next(dataloader_iter)
              (
                  tokens,
                  types,
                  sentence_order,
                  loss_mask,
                  lm_labels,
                  padding_mask,
              ) = (
                  batch["text"].cuda(non_blocking=True),
                  batch["types"].cuda(non_blocking=True),
                  batch["is_random"].cuda(non_blocking=True),
                  batch["loss_mask"].cuda(non_blocking=True),
                  batch["labels"].cuda(non_blocking=True),
                  batch["padding_mask"].cuda(non_blocking=True),
              )
          else:
              batch = next(dataloader_iter)
              if parallel_state.is_pipeline_first_stage():
                  tokens = batch["text"].cuda(non_blocking=True)
                  types = batch["types"].cuda(non_blocking=True)
                  sentence_order = batch["is_random"].cuda(
                      non_blocking=True
                  )
                  padding_mask = batch["padding_mask"].cuda(
                      non_blocking=True
                  )
                  loss_mask, lm_labels = None, None
              elif parallel_state.is_pipeline_last_stage():
                  loss_mask = batch["loss_mask"].cuda(non_blocking=True)
                  lm_labels = batch["labels"].cuda(non_blocking=True)
                  sentence_order = batch["is_random"].cuda(
                      non_blocking=True
                  )
                  padding_mask = batch["padding_mask"].cuda(
                      non_blocking=True
                  )
                  tokens, types = None, None
              else:
                  padding_mask = batch["padding_mask"].cuda(
                      non_blocking=True
                  )
                  sentence_order = batch["is_random"].cuda(
                      non_blocking=True
                  )
                  tokens, types, loss_mask, lm_labels = (
                      None,
                      None,
                      None,
                      None,
                  )
  
          if not self.cfg.bert_binary_head:
              types = None
  
          forward_args = {
              "input_ids": tokens,
              "attention_mask": padding_mask,
              "lm_labels": lm_labels,
          }
  
          if not self.mcore_bert:
              forward_args[
                  "checkpoint_activations_all_layers"
              ] = checkpoint_activations_all_layers
              forward_args["model"] = model
              forward_args["token_type_ids"] = types
          else:
              forward_args["tokentype_ids"] = types
  
          output_tensor = None
          if self.mcore_bert:
              output_tensor = model(**forward_args)
          else:
              output_tensor = self.forward(**forward_args)
  
          def loss_func(output_tensor):
              loss_dict = self.loss_func(
                  loss_mask, sentence_order, output_tensor
              )
              if "sop loss" in loss_dict:
                  lm_loss = loss_dict["lm loss"]
                  sop_loss = loss_dict["sop loss"]
                  loss = lm_loss + sop_loss
                  reduced_loss = (
                      average_losses_across_data_parallel_group(
                          [loss, lm_loss, sop_loss]
                      )
                  )
              else:
                  lm_loss = loss_dict["lm loss"]
                  loss = lm_loss
                  reduced_loss = (
                      average_losses_across_data_parallel_group(
                          [loss, lm_loss]
                      )
                  )
              return loss, {"loss": reduced_loss}
  
          return output_tensor, loss_func
  
      return fwd_output_and_loss_func
  
imports: import itertools
  from typing import Any, Dict, List, Optional
  import torch
  import torch.nn.functional
  from omegaconf.dictconfig import DictConfig
  from pytorch_lightning.trainer.trainer import Trainer
  from nemo.collections.nlp.data.language_modeling.megatron import dataset_utils
  from nemo.collections.nlp.data.language_modeling.megatron.data_samplers import MegatronPretrainingRandomSampler, MegatronPretrainingSampler
  from nemo.collections.nlp.models.language_modeling.megatron.bert_model import BertModel
  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.utils import ApexGuardDefaults, average_losses_across_data_parallel_group, get_params_for_weight_decay_optimization
  from nemo.collections.nlp.parts.nlp_overrides import GradScaler
  from nemo.collections.nlp.parts.utils_funcs import get_last_rank
  from nemo.core.classes.common import PretrainedModelInfo
  from nemo.core.neural_types import ChannelType, MaskType, NeuralType
  from nemo.utils import AppState, logging
label_desc: Use image node_id 4 for calling the MegatronBertModel obj's underlying member method code with example usage: obj.get_forward_output_and_loss_func() and returns: fwd_output_and_loss_func, output_tensor,  loss_func, loss,  dict