| """Export a weights-only init from a training bundle: live state overlaid with the | |
| EMA shadows (the eval-time weights), _orig_mod stripped — the format load_checkpoint | |
| expects for warm-starting a new run.""" | |
| import sys, torch | |
| bundle, out = sys.argv[1], sys.argv[2] | |
| d = torch.load(bundle, map_location="cpu", weights_only=False) | |
| strip = lambda sd: {k.replace("_orig_mod.", ""): v for k, v in sd.items()} | |
| sd = strip(d["model"]) | |
| sd.update(strip(d["ema"])) | |
| torch.save(sd, out) | |
| print(f"{out}: {len(sd)} keys from {bundle}") | |