megatron.checkpointing.load_biencoder_checkpoint

megatron.checkpointing.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