def train_epoch(model, loader, optimizer, device):
    model.train()
    total_loss = 0.0
    for xb, yb in loader:
        xb, yb = xb.to(device), yb.to(device)
        optimizer.zero_grad()
        logits = model(xb)
        loss = F.cross_entropy(logits, yb)
        loss.backward()
        optimizer.step()
        total_loss += loss.item() * xb.size(0)
    return total_loss / len(loader.dataset)


def evaluate(model, loader, device):
    model.eval()
    correct = 0
    with torch.no_grad():
        for xb, yb in loader:
            xb, yb = xb.to(device), yb.to(device)
            preds = model(xb).argmax(dim=1)
            correct += (preds == yb).sum().item()
    return correct / len(loader.dataset)


if __name__ == "__main__":
    device = "cuda" if torch.cuda.is_available() else "cpu"
    tfm = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.5,), (0.5,)),
    ])
    # Fashion-MNIST: an open, public-domain-licensed drop-in
    # replacement for MNIST with 10 clothing categories.
    train_set = datasets.FashionMNIST(
        root="./data", train=True, download=True, transform=tfm
    )
    test_set = datasets.FashionMNIST(
        root="./data", train=False, download=True, transform=tfm
    )
    train_loader = DataLoader(train_set, batch_size=128, shuffle=True)
    test_loader = DataLoader(test_set, batch_size=256)

    model = SimpleCNN(num_classes=10).to(device)
    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

    for epoch in range(3):
        train_loss = train_epoch(model, train_loader, optimizer, device)
        acc = evaluate(model, test_loader, device)
        print(f"epoch {epoch}  loss {train_loss:.4f}  test_acc {acc:.3f}")
