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 forward(
      self,
      input_ids,
      attention_mask,
      token_type_ids,
      lm_labels=None,
      checkpoint_activations_all_layers=None,
      model=None,
  ):
      if model is None:
          model = self.model
      output_tensor = model(
          input_ids,
          attention_mask,
          token_type_ids=token_type_ids,
          lm_labels=lm_labels,
          checkpoint_activations_all_layers=checkpoint_activations_all_layers,
      )
      if parallel_state.is_pipeline_last_stage():
          # Return the output tensor of encoder and transpose from [seq_len, batch, hidden] to [batch, seq_len, hidden]
          if torch.is_tensor(output_tensor):
              output_tensor = output_tensor.transpose(1, 0).contiguous()
          else:
              lm_loss_, sop_logits = output_tensor
  
              lm_loss_ = lm_loss_.transpose(1, 0).contiguous()
              if sop_logits is not None:
                  sop_logits = sop_logits.transpose(1, 0).contiguous()
              output_tensor = (lm_loss_, sop_logits)
  
      return output_tensor
  
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 5 for calling the MegatronBertModel obj's underlying member method code with example usage: obj.forward(input_ids, attention_mask, token_type_ids, lm_labels, checkpoint_activations_all_layers, model) and returns: output_tensor