if __name__ == "__main__":
    torch.manual_seed(0)
    width, depth, batch = 256, 40, 64
    model = CheckpointedMLP(width, depth)
    x = torch.randn(batch, width, requires_grad=True)

    def run_and_measure(use_checkpoint):
        out = model(x, use_checkpoint=use_checkpoint)
        loss = out.sum()
        loss.backward()
        if torch.cuda.is_available():
            peak = torch.cuda.max_memory_allocated() / 1e6
            torch.cuda.reset_peak_memory_stats()
            return peak
        return None

    peak_normal = run_and_measure(use_checkpoint=False)
    peak_checkpointed = run_and_measure(use_checkpoint=True)
    if peak_normal is not None:
        print(f"Peak memory (normal):       {peak_normal:.1f} MB")
        print(f"Peak memory (checkpointed): {peak_checkpointed:.1f} MB")
    else:
        print(
            "No GPU available; rerun on CUDA to see the memory "
            "reduction from activation checkpointing directly."
        )

    # gradient accumulation demo: 4 micro-batches of 16 == 1
    # batch of 64, but each micro-batch only needs 16-example
    # memory at any one time.
    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
    micro_batches = [
        (torch.randn(16, width), torch.zeros(16, dtype=torch.long))
        for _ in range(4)
    ]
    # NOTE: this model outputs `width`-dim vectors, not class
    # logits; in a real classifier, model output would already
    # be sized to the number of classes.
