def train_step(self, graph, feat, temperature, **kwargs):
        r"""Compute the loss of the explanation network for graph classification

        Parameters
        ----------
        graph : DGLGraph
            Input batched homogeneous graph.
        feat : Tensor
            The input feature of shape :math:`(N, D)`. :math:`N` is the
            number of nodes, and :math:`D` is the feature size.
        temperature : float
            The temperature parameter fed to the sampling procedure.
        kwargs : dict
            Additional arguments passed to the GNN model.

        Returns
        -------
        Tensor
            A scalar tensor representing the loss.
        """
        assert (
            self.graph_explanation
        ), '"explain_graph" must be True when initializing the module.'

        self.model = self.model.to(graph.device)
        self.elayers = self.elayers.to(graph.device)

        pred = self.model(graph, feat, embed=False, **kwargs)
        pred = pred.argmax(-1).data

        prob, _ = self.explain_graph(
            graph, feat, temperature, training=True, **kwargs
        )

        loss = self.loss(prob, pred)
        return loss
Compute the loss of the explanation network for graph classification

        Parameters
        ----------
        graph : DGLGraph
            Input batched homogeneous graph.
        feat : Tensor
            The input feature of shape :math:`(N, D)`. :math:`N` is the
            number of nodes, and :math:`D` is the feature size.
        temperature : float
            The temperature parameter fed to the sampling procedure.
        kwargs : dict
            Additional arguments passed to the GNN model.

        Returns
        -------
        Tensor
            A scalar tensor representing the loss.
        
python/dgl/nn/pytorch/explain/pgexplainer.py