"""Average the weights of two epoch checkpoints (the z2 model is the 50/50 average of epochs 0 and 1).

    python average_checkpoints.py runs/z2/epoch0.pt runs/z2/epoch1.pt --out runs/z2/avg.pt
"""
import argparse

import torch

ap = argparse.ArgumentParser()
ap.add_argument("a")
ap.add_argument("b")
ap.add_argument("--out", required=True)
args = ap.parse_args()
a = torch.load(args.a, map_location="cpu", weights_only=False)
b = torch.load(args.b, map_location="cpu", weights_only=False)
assert a["config"] == b["config"] and a["model"].keys() == b["model"].keys()
avg = {k: (0.5 * a["model"][k].float() + 0.5 * b["model"][k].float()).to(a["model"][k].dtype)
        if a["model"][k].is_floating_point() else a["model"][k] for k in a["model"]}
torch.save({"model": avg, "config": a["config"]}, args.out)