Recovering graph indices from a batch

[oiao]

Consider the following dataset:

import numpy as np
import torch
from torch_geometric.data import Data, InMemoryDataset
from torch_geometric.loader import DataLoader

datalist = []
for _ in range(100):
    size = np.random.randint(10)
    datalist.append( Data(x=torch.randn(size, 1)) )
dataset = InMemoryDataset()
dataset.data, dataset.slices = dataset.collate(datalist)

If I batch the dataset with

loader = DataLoader(dataset, batch_size=5, shuffle=True)
for batch in loader:
    break

Is there any way for me to recover the index of each graph in batch, such that after going through all of the batches I could easily restore the original sequence of individual graphs, as given by datalist? I am aware that I could add an additional index attribute to my dataset, but is there a way to solve this out-of-the box? Maybe some internal attribute of batch that I am missing?

Thanks a lot!

[rusty1s]

AFAIK, the only way to do it is to create such an index tensor. PyTorch does not seem to provide such a functionality otherwise.

[oiao]

Thanks a lot for the answer.