How to get original data index from each batch

[RX28666]

I would like to know if there is a way to retrieve the indices of the data points in a minibatch in PyG?

eg:

dataset = TUDataset(root='/tmp/ENZYMES', name='ENZYMES', use_node_attr=True)

loader = DataLoader(dataset, batch_size=32, shuffle=True)

how can I get the the indices (in original ENZYMES dataset) of those 32 datas in each batch?
Thanks.

[rusty1s]

You can add them to the dataset before-hand:

data_list = []
for i, data in enumerate(dataset)
    data.graph_idx = i
    data_list.append(data)

loader = DataLoader(data_list, batch_size=32, shuffle=True)

batch = next(iter(loader))
print(batch.graph_idx)

[RX28666]

It works. Thanks so much.