class_docstr: Bert-based module for Inverse Cloze task.
class_docstr_tok: ["Bert-based","module","for","Inverse","Cloze","task","."]
docstr: Save dict with state dicts of each of the models.
docstr_tok: ["Save","dict","with","state","dicts","of","each","of","the","models","."]
code: def state_dict_for_save_checkpoint(self, prefix="", keep_vars=False):
      """Save dict with state dicts of each of the models."""
      state_dict_ = {}
      if self.use_query_model:
          state_dict_[
              self._query_key
          ] = self.query_model.state_dict_for_save_checkpoint(
              prefix=prefix, keep_vars=keep_vars
          )
  
      if self.use_block_model:
          state_dict_[
              self._block_key
          ] = self.block_model.state_dict_for_save_checkpoint(
              prefix=prefix, keep_vars=keep_vars
          )
  
      return state_dict_
  
imports: import os
  import torch
  from megatron import get_args, print_rank_0
  from megatron.checkpointing import get_checkpoint_tracker_filename, get_checkpoint_name
  from megatron.model import BertModel
  from .module import MegatronModule
  from megatron.core import mpu
  from megatron.model.enums import AttnMaskType
  from megatron.model.utils import get_linear_layer
  from megatron.model.utils import init_method_normal
  from megatron.model.language_model import get_language_model
  from megatron.model.utils import scaled_init_method_normal
  from megatron.model.bert_model import bert_extended_attention_mask, bert_position_ids
label_desc: Use image node_id 5 for calling the ICTBertModel obj's underlying member method code with example usage: obj.state_dict_for_save_checkpoint(prefix, keep_vars) and returns: state_dict_