Indices graphs in original dataset

Hi!

When generating a batch from a torch_geometric loader object while training, is it possible to refer back to the indices in the original dataset? i.e., let’s say my loader is based on the graph dataset [g_0, g_1, g_2, g_3, g_4, g_5] and my batch is based on graphs [g_3, g_1], is it possible to retrieve indices (3, 1)?

Thanks!

I’m not sure how torch_geometric implements their datasets, but assuming they are similar to torch.utils.data.Dataset (or the same), you could return the index in the __getitem__ and retrieve it in the DataLoader.

1 Like

Thanks! This worked :slight_smile: To be more precise (for someone else perhaps encountering this), I altered the get function (see torch_geometric.data.in_memory_dataset — pytorch_geometric 2.0.5 documentation) to return the index. Thanks again!