| """usage: python soup.py OUT IN1 IN2 ... (uniform weight soup; int buffers from first)""" | |
| import sys, torch | |
| sds = [torch.load(p, map_location="cpu", weights_only=True)["state_dict"] for p in sys.argv[2:]] | |
| s = {k: ((sum(d[k].float() for d in sds) / len(sds)).to(sds[0][k].dtype) if sds[0][k].is_floating_point() else sds[0][k]) for k in sds[0]} | |
| torch.save({"state_dict": s}, sys.argv[1]); print("soup", sys.argv[1], len(sds), flush=True) | |