blog_9.py

import torch_geometric.transforms as T

class CoraDataset(LightningDataModule):

    NAME = "cora"
    
    ...

    def gather_data(self, batch, batch_nb):
        """
        This function will select features using node_idx
        and convert it into a Namedtuple
        """
        usual_keys = ["x", "edge_index", "edge_attr", "batch"]
        Batch: TensorBatch = namedtuple("Batch", usual_keys)
        return (
            Batch(
                self.data.x[batch[1]],
                [e.edge_index for e in batch[2]],
                None, # Not used as we don't have any attr for the edge
                None # tensor associated its node to its orginal batch index
            ),
            self.data.y[batch[1]],
        )