1from __future__ import annotations23from pathlib import Path45import torch678def save(path: Path, model, optimizer, scheduler, step: int, best: float, extra: dict | None = None) -> None:9 path.parent.mkdir(parents=True, exist_ok=True)10 tmp = path.with_suffix(".tmp")11 torch.save({"model": model.state_dict(), "optimizer": optimizer.state_dict() if optimizer else None,12 "scheduler": scheduler.state_dict() if scheduler else None, "step": step, "best": best,13 "cfg": model.cfg.__dict__, "extra": extra or {}}, tmp)14 tmp.replace(path)151617def load(path: Path, model, optimizer=None, scheduler=None, map_location="cpu") -> dict:18 ck = torch.load(path, map_location=map_location, weights_only=False)19 model.load_state_dict(ck["model"])20 if optimizer and ck.get("optimizer"):21 optimizer.load_state_dict(ck["optimizer"])22 if scheduler and ck.get("scheduler"):23 scheduler.load_state_dict(ck["scheduler"])24 return ck25