Function load_checkpoint_for_inference

generate.py:61–64

(ckpt_path, model, device="cpu")

def load_checkpoint_for_inference(ckpt_path, model, device="cpu"):
    checkpoint = torch.load(ckpt_path, map_location=device)
    raw_model = model.module if hasattr(model, "module") else model
    raw_model.load_state_dict(checkpoint["model_state_dict"], strict=False)


@torch.no_grad()