def sample(self, g, partition_ids):  # pylint: disable=arguments-differ
        """Sampling function.

        Parameters
        ----------
        g : DGLGraph
            The graph to sample from.
        partition_ids : Tensor
            A 1-D integer tensor of partition IDs.

        Returns
        -------
        DGLGraph
            The sampled subgraph.
        """
        node_ids = F.cat(
            [
                self.partition_node_ids[
                    self.partition_offset[i] : self.partition_offset[i + 1]
                ]
                for i in F.asnumpy(partition_ids)
            ],
            0,
        )
        sg = g.subgraph(
            node_ids, relabel_nodes=True, output_device=self.output_device
        )
        set_node_lazy_features(sg, self.prefetch_ndata)
        set_edge_lazy_features(sg, self.prefetch_edata)
        return sg
Sampling function.

        Parameters
        ----------
        g : DGLGraph
            The graph to sample from.
        partition_ids : Tensor
            A 1-D integer tensor of partition IDs.

        Returns
        -------
        DGLGraph
            The sampled subgraph.
        
python/dgl/dataloading/cluster_gcn.py