class_docstr: BERT-based encoder for queries or contexts used for
  learned information retrieval.
class_docstr_tok: ["BERT-based","encoder","for","queries","or","contexts","used","for","learned","information","retrieval","."]
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, 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(
          prefix=prefix, keep_vars=keep_vars
      )
  
      if self.biencoder_projection_dim > 0:
          state_dict_[
              self._projection_enc_key
          ] = self.projection_enc.state_dict(
              prefix=prefix, keep_vars=keep_vars
          )
  
      return state_dict_
  
imports: import os
  import torch
  import sys
  from megatron import get_args, print_rank_0, get_tokenizer
  from megatron.core import mpu
  from megatron.checkpointing import fix_query_key_value_ordering
  from megatron.checkpointing import get_checkpoint_tracker_filename
  from megatron.checkpointing import get_checkpoint_name
  from megatron.model.bert_model import bert_position_ids
  from megatron.model.enums import AttnMaskType
  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 3 for calling the PretrainedBertModel obj's underlying member method code with example usage: obj.state_dict_for_save_checkpoint(prefix, keep_vars) and returns: state_dict_