class_docstr: Megatron Bert pretraining.
  Model returns [batch, seq, hidden] shape
class_docstr_tok: ["Megatron","Bert","pretraining",".","Model","returns","[","batch",",","seq",",","hidden","]","shape"]
docstr: Generates input examples for tracing etc.
  Returns:
      A tuple of input examples.
code: def input_example(self, max_batch=1, max_dim=256):
      """
      Generates input examples for tracing etc.
      Returns:
          A tuple of input examples.
      """
      sample = next(self.parameters())
      sz = (max_batch, max_dim)
      input_ids = torch.randint(
          low=0, high=2048, size=sz, device=sample.device
      )
      token_type_ids = torch.randint(
          low=0, high=1, size=sz, device=sample.device
      )
      attention_mask = torch.randint(
          low=0, high=1, size=sz, device=sample.device
      )
      input_dict = {
          "input_ids": input_ids,
          "attention_mask": attention_mask,
          "token_type_ids": token_type_ids,
      }
      return tuple([input_dict])
  
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 32 for calling the MegatronBertModel obj's underlying member method code with example usage: obj.input_example(max_batch, max_dim) and returns: tuple