def forward(self, block=None):
        """Forward computation

        Parameters
        ----------
        block : DGLGraph, optional
            If not specified, directly return the full graph with embeddings stored in
            :attr:`embed_name`. Otherwise, extract and store the embeddings to the block
            graph and return.

        Returns
        -------
        DGLGraph
            The block graph fed with embeddings.
        """
        embeds = {}
        for ntype in block.ntypes:
            if self.node_feats[ntype] is None:
                embeds[ntype] = self.node_embeds[ntype](block.nodes(ntype)).to(
                    self.device
                )
            else:
                embeds[ntype] = (
                    self.node_feats[ntype][block.nodes(ntype)].to(self.device)
                    @ self.embeds[ntype]
                )
        return embeds
Forward computation

        Parameters
        ----------
        block : DGLGraph, optional
            If not specified, directly return the full graph with embeddings stored in
            :attr:`embed_name`. Otherwise, extract and store the embeddings to the block
            graph and return.

        Returns
        -------
        DGLGraph
            The block graph fed with embeddings.
        
benchmarks/benchmarks/model_speed/bench_rgcn_hetero_ns.py