docstr: Customized load.
docstr_tok: ["Customized","load","."]
code: def load_state_dict(self, state_dict, strict=True):
      """Customized load."""
  
      self.language_model.load_state_dict(
          state_dict[self._language_model_key], strict=strict
      )
      if self.post_process:
          if self._classification_head_key in state_dict:
              self.classification_head.load_state_dict(
                  state_dict[self._classification_head_key],
                  strict=strict,
              )
          else:
              print_rank_last(
                  "***WARNING*** could not find {} in the checkpoint, "
                  "initializing to random".format(
                      self._classification_head_key
                  )
              )
  
file_path: megatron-lm/megatron/model/classification.py
imports: import torch
  from megatron import get_args, print_rank_last
  from megatron.model.enums import AttnMaskType
  from megatron.model.bert_model import bert_extended_attention_mask, bert_position_ids
  from megatron.model.language_model import get_language_model
  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 5 for calling the Classification obj's underlying member method code with example usage: obj.load_state_dict(state_dict, strict) without return types