class_docstr: Bert-based module for Biencoder model.
class_docstr_tok: ["Bert-based","module","for","Biencoder","model","."]
docstr: Load the state dicts of each of the models
docstr_tok: ["Load","the","state","dicts","of","each","of","the","models"]
code: def load_state_dict(self, state_dict, strict=True):
      """Load the state dicts of each of the models"""
      if self.biencoder_shared_query_context_model:
          print_rank_0("Loading shared query-context model")
          self.model.load_state_dict(
              state_dict[self._model_key], strict=strict
          )
      else:
          if self.use_query_model:
              print_rank_0("Loading query model")
              self.query_model.load_state_dict(
                  state_dict[self._query_key], strict=strict
              )
  
          if self.use_context_model:
              print_rank_0("Loading context model")
              self.context_model.load_state_dict(
                  state_dict[self._context_key], strict=strict
              )
  
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 6 for calling the BiEncoderModel obj's underlying member method code with example usage: obj.load_state_dict(state_dict, strict) without return types