def forward(self, g, x, *args):
        r"""Apply the GNN module with grouped reversible residual connection.

        Parameters
        ----------
        g : DGLGraph
            The graph.
        x : torch.Tensor
            The input feature of shape :math:`(N, D_{in})`, where :math:`D_{in}` is size
            of input feature, :math:`N` is the number of nodes.
        args
            Additional arguments to pass to :attr:`gnn_module`.

        Returns
        -------
        torch.Tensor
            The output feature of shape :math:`(N, D_{in})`.
        """
        args = (g, x) + args
        y = InvertibleCheckpoint.apply(
            self._forward,
            self._inverse,
            len(args),
            *(args + tuple([p for p in self.parameters() if p.requires_grad]))
        )

        return y
Apply the GNN module with grouped reversible residual connection.

        Parameters
        ----------
        g : DGLGraph
            The graph.
        x : torch.Tensor
            The input feature of shape :math:`(N, D_{in})`, where :math:`D_{in}` is size
            of input feature, :math:`N` is the number of nodes.
        args
            Additional arguments to pass to :attr:`gnn_module`.

        Returns
        -------
        torch.Tensor
            The output feature of shape :math:`(N, D_{in})`.
        
python/dgl/nn/pytorch/conv/grouprevres.py