Class CheckpointedNet

tests/test_linear4bit.py:483–503

    torch.compiler.reset()

    class CheckpointedNet(torch.nn.Module):
        def __init__(self):
            super().__init__()
            self.layers = torch.nn.ModuleList(
                [
                    bnb.nn.Linear4bit(
                        dim,
                        dim,
                        bias=False,
                        compute_dtype=compute_dtype,
                        compress_statistics=compress_statistics,
                        quant_type=quant_type,
                    )
                    for _ in range(4)
                ]
            )

        def forward(self, x):
            for layer in self.layers:
                x = torch.utils.checkpoint.checkpoint(layer, x, use_reentrant=False)
            return x

    net = CheckpointedNet().to(device)