import torch
from torch import nn


torch.manual_seed(7)
checkpoint_path = "training-checkpoint.tar"


def build_model():
    return nn.Sequential(
        nn.Linear(3, 4),
        nn.ReLU(),
        nn.Linear(4, 1),
    )


x = torch.tensor(
    [
        [0.1, 0.2, 0.3],
        [0.4, 0.5, 0.6],
    ],
    dtype=torch.float32,
)
y = torch.tensor(
    [
        [0.6],
        [1.5],
    ],
    dtype=torch.float32,
)
loss_fn = nn.MSELoss()


def train_step(model, optimizer):
    model.train()
    optimizer.zero_grad()
    loss = loss_fn(model(x), y)
    loss.backward()
    optimizer.step()
    return loss.item()


model = build_model()
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)

loss_before_save = train_step(model, optimizer)
torch.save(
    {
        "epoch": 1,
        "model_state_dict": model.state_dict(),
        "optimizer_state_dict": optimizer.state_dict(),
        "loss": loss_before_save,
    },
    checkpoint_path,
)
print(f"saved epoch=1 loss={loss_before_save:.4f}")

restored_model = build_model()
restored_optimizer = torch.optim.Adam(restored_model.parameters(), lr=0.01)
checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=True)

restored_model.load_state_dict(checkpoint["model_state_dict"])
restored_optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
restored_model.train()

state_match = all(
    torch.equal(model.state_dict()[name], restored_model.state_dict()[name])
    for name in model.state_dict()
)
optimizer_state_entries = len(restored_optimizer.state_dict()["state"])
next_epoch = checkpoint["epoch"] + 1
resumed_loss = train_step(restored_model, restored_optimizer)

print(f"loaded keys={', '.join(checkpoint.keys())}")
print(f"state_match={state_match}")
print(f"optimizer_state_entries={optimizer_state_entries}")
print(f"resume_epoch={next_epoch} previous_loss={checkpoint['loss']:.4f}")
print(f"resumed_loss={resumed_loss:.4f}")