Function pre_load

modelscope/utils/nlp/load_checkpoint.py:65–71

(mp_rank, load_dir, tag='')

def pre_load(mp_rank, load_dir, tag=''):
    load_path = _get_ckpt_name(mp_rank, load_dir, tag)
    checkpoint = torch.load(
        load_path,
        map_location=lambda storage, loc: storage,
        weights_only=True)
    return checkpoint['module'] if 'module' in checkpoint else checkpoint


def _load_checkpoint(model,