class_docstr: Bert Language model.
class_docstr_tok: ["Bert","Language","model","."]
code: def forward(
      self,
      bert_model_input,
      attention_mask,
      tokentype_ids=None,
      lm_labels=None,
  ):
      extended_attention_mask = bert_extended_attention_mask(
          attention_mask
      )
      input_ids = bert_model_input
      position_ids = bert_position_ids(input_ids)
  
      lm_output = self.language_model(
          input_ids,
          position_ids,
          extended_attention_mask,
          tokentype_ids=tokentype_ids,
      )
  
      if self.post_process and self.add_binary_head:
          lm_output, pooled_output = lm_output
  
          # Return pooled output (e.g., when computing Bert embeddings).
          if self.return_embeddings:
              # Sum attention mask.
              embeddings = torch.transpose(lm_output, 0, 1)
              masks = torch.sum(attention_mask, dim=1)
  
              # Collect masked embeddings.
              output = torch.zeros(
                  size=(embeddings.shape[0], embeddings.shape[2]),
                  dtype=torch.float32,
                  device=torch.cuda.current_device(),
              )
              for i, (embedding, mask) in enumerate(
                  zip(embeddings, masks)
              ):
                  output[i, :] = torch.mean(
                      embedding[1 : mask - 1], dim=0
                  )
  
              return output
  
      else:
          pooled_output = None
  
      if self.post_process:
          return post_language_model_processing(
              lm_output,
              pooled_output,
              self.lm_head,
              self.binary_head,
              lm_labels,
              self.shared_embedding_or_output_weight(),
              self.fp16_lm_cross_entropy,
          )
      else:
          return lm_output
  
imports: import torch
  from megatron import get_args
  from megatron.core import tensor_parallel
  from megatron.model.enums import AttnMaskType
  from megatron.model.language_model import parallel_lm_logits
  from megatron.model.language_model import get_language_model
  from megatron.model.utils import get_norm
  from megatron.model.utils import openai_gelu, erf_gelu
  from megatron.model.utils import get_linear_layer
  from megatron.model.utils import init_method_normal
  from megatron.model.utils import scaled_init_method_normal
  from .module import MegatronModule
label_desc: Use image node_id 3 for calling the BertModel obj's underlying member method code with example usage: obj.forward(bert_model_input, attention_mask, tokentype_ids, lm_labels) and returns: post_language_model_processing, lm_output, output