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 loss_func(self, loss_mask, sentence_order, output_tensor):
      lm_loss_, sop_logits = output_tensor
  
      lm_loss_ = lm_loss_.float()
      loss_mask = loss_mask.float()
  
      # Sometimes when the number of tokens is very small, none of the tokens get masked for prediction. In that case loss mask is all zeros
      # i.e Happens when the entire batch is masked out (Practically when MBS=1 or 2, and the number of tokens in each batch is < 7 )
      if loss_mask.sum() == 0:
          lm_loss = torch.sum(lm_loss_.view(-1)) * 0.0
      else:
          lm_loss = (
              torch.sum(lm_loss_.view(-1) * loss_mask.reshape(-1))
              / loss_mask.sum()
          )
  
      if sop_logits is not None:
          sop_loss = F.cross_entropy(
              sop_logits.view(-1, 2).float(),
              sentence_order.view(-1),
              ignore_index=-1,
          )
          sop_loss = sop_loss.float()
          return {"lm loss": lm_loss, "sop loss": sop_loss}
          # loss = lm_loss + sop_loss
          # averaged_losses = average_losses_across_data_parallel_group(
          #     [lm_loss, sop_loss])
          # return loss, {'lm loss': averaged_losses[0],
          #               'sop loss': averaged_losses[1]}
  
      else:
          return {"lm loss": lm_loss}
  
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 12 for calling the MegatronBertModel obj's underlying member method code with example usage: obj.loss_func(loss_mask, sentence_order, output_tensor) and returns: dict, dict