class_docstr: Bert Language model.
docstr: For easy load when model is combined with other heads,
  add an extra key.
docstr_tok: ["For","easy","load","when","model","is","combined","with","other","heads",",","add","an","extra","key","."]
code: def state_dict_for_save_checkpoint(
      self, destination=None, prefix="", keep_vars=False
  ):
      """For easy load when model is combined with other heads,
      add an extra key."""
  
      state_dict_ = {}
      state_dict_[
          self._language_model_key
      ] = self.language_model.state_dict_for_save_checkpoint(
          destination, prefix, keep_vars
      )
      if self.post_process:
          state_dict_[
              self._lm_head_key
          ] = self.lm_head.state_dict_for_save_checkpoint(
              destination, prefix, keep_vars
          )
      if self.post_process and self.add_binary_head:
          state_dict_[
              self._binary_head_key
          ] = self.binary_head.state_dict(
              destination, prefix, keep_vars
          )
      # Save word_embeddings.
      if self.post_process and not self.pre_process:
          state_dict_[
              self._word_embeddings_for_head_key
          ] = self.word_embeddings.state_dict(
              destination, prefix, keep_vars
          )
      return state_dict_
  
imports: import contextlib
  import torch
  from apex.transformer import tensor_parallel
  from apex.transformer.enums import AttnMaskType
  from apex.transformer.enums import ModelType
  from apex.transformer.layers import FusedLayerNorm
  from apex.transformer.testing.global_vars import get_args
  from apex.transformer.testing.standalone_transformer_lm import MegatronModule, get_language_model, get_linear_layer, init_method_normal, scaled_init_method_normal, parallel_lm_logits
label_desc: Use image node_id 4 for calling the BertModel obj's underlying member method code with example usage: obj.state_dict_for_save_checkpoint(destination, prefix, keep_vars) and returns: state_dict_