Function _load_checkpoint

open_instruct/model_utils.py:250–258

(path: str, dev: torch.device)

    """

    def _load_checkpoint(path: str, dev: torch.device):
        state_dict = torch.load(path, map_location=dev)
        if hasattr(model, "module"):
            # Needed if wrapped by DeepSpeed.
            model.module.load_state_dict(state_dict)
        else:
            # If a vanilla HF model.
            model.load_state_dict(state_dict)
        logger.info(f"{rank=}: Loaded checkpoint from {path}")

    if not throw_on_error:
        try: