model/utils.py

from pathlib import Path


def latest_ckpt(d: str | Path):
    """Latest checkpoint: prefer model_final.pt, else max numeric step."""
    d = Path(d)
    final = d / "model_final.pt"
    if final.exists():
        return final
    ckpts = [p for p in d.glob("model_*.pt") if p.stem.split("_")[-1].isdigit()]
    return max(ckpts, key=lambda p: int(p.stem.split("_")[-1])) if ckpts else None