code: def load_biencoder_checkpoint(
      model,
      only_query_model=False,
      only_context_model=False,
      custom_load_path=None,
  ):
      """
      selectively load retrieval models for indexing/retrieving
      from saved checkpoints
      """
  
      args = get_args()
  
      model = unwrap_model(model)
  
      load_path = (
          custom_load_path
          if custom_load_path is not None
          else args.load
      )
  
      tracker_filename = get_checkpoint_tracker_filename(load_path)
      with open(tracker_filename, "r") as f:
          iteration = int(f.read().strip())
  
      checkpoint_name = get_checkpoint_name(
          load_path,
          iteration,
          args.use_distributed_optimizer,
          release=False,
      )
  
      if mpu.get_data_parallel_rank() == 0:
          print(
              "global rank {} is loading checkpoint {}".format(
                  torch.distributed.get_rank(), checkpoint_name
              )
          )
  
      state_dict = torch.load(checkpoint_name, map_location="cpu")
      ret_state_dict = state_dict["model"]
  
      if only_query_model:
          ret_state_dict.pop("context_model")
      if only_context_model:
          ret_state_dict.pop("query_model")
  
      assert len(model) == 1
      model[0].load_state_dict(ret_state_dict)
      torch.distributed.barrier()
  
      if mpu.get_data_parallel_rank() == 0:
          print(" successfully loaded {}".format(checkpoint_name))
  
      return model
  
imports: import os
  import random
  import sys
  import numpy
  import torch
  from megatron import update_num_microbatches
  from megatron.core import mpu, tensor_parallel
  from .global_vars import get_args
  from .utils import unwrap_model, print_rank_0
label_desc: Use image node_id 17 for calling a global function with example usage: load_biencoder_checkpoint(model, only_query_model, only_context_model, custom_load_path) and returns: model