def forward(
            ctx,
            x,
            layers,
            buffer_layers,  # List of layer ids for int activation to buffer
    ):
        """Reversible Forward pass.

        Any intermediate activations from `buffer_layers` are cached in ctx for
        forward pass. This is not necessary for standard usecases. Each
        reversible layer implements its own forward pass logic.
        """
        buffer_layers.sort()
        x1, x2 = torch.chunk(x, 2, dim=-1)
        intermediate = []

        for layer in layers:
            x1, x2 = layer(x1, x2)
            if layer.layer_id in buffer_layers:
                intermediate.extend([x1.detach(), x2.detach()])

        if len(buffer_layers) == 0:
            all_tensors = [x1.detach(), x2.detach()]
        else:
            intermediate = [torch.LongTensor(buffer_layers), *intermediate]
            all_tensors = [x1.detach(), x2.detach(), *intermediate]

        ctx.save_for_backward(*all_tensors)
        ctx.layers = layers

        return torch.cat([x1, x2], dim=-1)
Reversible Forward pass.

        Any intermediate activations from `buffer_layers` are cached in ctx for
        forward pass. This is not necessary for standard usecases. Each
        reversible layer implements its own forward pass logic.
        
mmpretrain/models/backbones/revvit.py