checkpoint_path = logdir / "checkpoints" / f"{checkpoint_name}.pth"
print("Load config")
config: Dict[str, dict] = load_config(config_path)
runner_params = config.get("runner_params", {}) or {}
// Get expdir name
config_expdir = Path(config["args"]["expdir"])
// We will use copy of expdir from logs for reproducibility
expdir = Path(logdir) / "code" / config_expdir.name
print("Import experiment and runner from logdir")
experiment_fn, runner_fn = import_experiment_and_runner(expdir)
experiment: ConfigExperiment = experiment_fn(config)
print(f"Load model state from checkpoints/{checkpoint_name}.pth")
if stage is None:
stage = list(experiment.stages)[0]
model = experiment.get_model(stage)
checkpoint = load_checkpoint(checkpoint_path)
unpack_checkpoint(checkpoint, model=model)
runner: runner_fn = runner_fn(**runner_params)
runner.model, runner.device = model, device
if loader is None:
loader = 0
After Change
runner.model = model_dump
return result
logger.info("Tracing is running...")
traced_model = trace_model(
model=model,
predict_fn=predict_fn,
batch=batch,