def backward(ctx, dx):
        """Reversible Backward pass.

        Any intermediate activations from `buffer_layers` are recovered from
        ctx. Each layer implements its own loic for backward pass (both
        activation recomputation and grad calculation).
        """
        d_x1, d_x2 = torch.chunk(dx, 2, dim=-1)
        # retrieve params from ctx for backward
        x1, x2, *int_tensors = ctx.saved_tensors
        # no buffering
        if len(int_tensors) != 0:
            buffer_layers = int_tensors[0].tolist()
        else:
            buffer_layers = []

        layers = ctx.layers

        for _, layer in enumerate(layers[::-1]):
            if layer.layer_id in buffer_layers:
                x1, x2, d_x1, d_x2 = layer.backward_pass(
                    y1=int_tensors[buffer_layers.index(layer.layer_id) * 2 +
                                   1],
                    y2=int_tensors[buffer_layers.index(layer.layer_id) * 2 +
                                   2],
                    d_y1=d_x1,
                    d_y2=d_x2,
                )
            else:
                x1, x2, d_x1, d_x2 = layer.backward_pass(
                    y1=x1,
                    y2=x2,
                    d_y1=d_x1,
                    d_y2=d_x2,
                )

        dx = torch.cat([d_x1, d_x2], dim=-1)

        del int_tensors
        del d_x1, d_x2, x1, x2

        return dx, None, None
Reversible Backward pass.

        Any intermediate activations from `buffer_layers` are recovered from
        ctx. Each layer implements its own loic for backward pass (both
        activation recomputation and grad calculation).
        
mmpretrain/models/backbones/revvit.py