"""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)