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