def train_vae(epochs=10, batch_size=128, latent_dim=20, lr=1e-3):
    tfm = transforms.Compose([transforms.ToTensor()])
    train_set = datasets.MNIST(
        root="./data", train=True, download=True, transform=tfm
    )
    loader = DataLoader(train_set, batch_size=batch_size, shuffle=True)

    model = VAE(latent_dim=latent_dim)
    optimizer = torch.optim.Adam(model.parameters(), lr=lr)

    for epoch in range(epochs):
        total_loss = 0.0
        for x, _ in loader:
            x = x.view(x.size(0), -1)
            x_hat, mu, logvar = model(x)
            loss, recon, kl = vae_loss(x_hat, x, mu, logvar)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            total_loss += loss.item()
        avg_loss = total_loss / len(train_set)
        print(f"epoch {epoch}  avg ELBO loss {avg_loss:.2f}")

    return model


@torch.no_grad()
def sample_vae(model, n_samples=16, latent_dim=20):
    z = torch.randn(n_samples, latent_dim)
    samples = model.decode(z)
    return samples.view(n_samples, 28, 28)


if __name__ == "__main__":
    torch.manual_seed(0)
    model = train_vae(epochs=5)
    samples = sample_vae(model, n_samples=16)
    # samples: (16, 28, 28) in [0, 1], ready to visualize as
    # a grid of generated MNIST-style digit images.
    print("Generated sample batch shape:", samples.shape)
