How to retrieve the sample indices of a mini-batch

One way to do this is to implement a subclass of torch.utils.data.Dataset that returns a triple (data, target, index) from its __getitem__ method. Then your loop would be:

for data, target, index in train_loader:
    ....
9 Likes