class_docstr: Masked LM head for Bert
  
  Arguments:
      config: TransformerConfig object
      mpu_vocab_size: model parallel size of vocabulary.
      parallel_output: whether output logits being distributed or not.
class_docstr_tok: ["Masked","LM","head","for","Bert","Arguments",":","config",":","TransformerConfig","object","mpu_vocab_size",":","model","parallel","size","of","vocabulary",".","parallel_output",":","whether","output","logits","being","distributed","or","not","."]
docstr: Customize load.
docstr_tok: ["Customize","load","."]
code: def load_state_dict(self, state_dict, strict=True):
      """Customize load."""
  
      # Handle renaming layernorm -> norm in component names
      state_dict_ = {}
      for key in state_dict.keys():
          newkey = key.replace("layernorm", "norm")
          state_dict_[newkey] = state_dict[key]
  
      super().load_state_dict(state_dict_, strict)
  
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 BertLMHead obj's underlying member method code with example usage: obj.load_state_dict(state_dict, strict) without return types