I am creating a dataloader from a Subset as follows:
dataset = torch.utils.data.TensorDataset(torch.rand(100))
train_subset, test_subset = torch.utils.data.random_split(dataset, [80, 20])
train_loader = torch.utils.data.DataLoader(train_subset, batch_size = 10)
test_loader = torch.utils.data.DataLoader(test_subset, batch_size = 10)

How can we get the i-th batch directly, without a for-loop? I am doing the following:
train_loader.dataset.dataset[train_loader.dataset.indices][0][list(train_loader.batch_sampler)[i]]
which looks ugly and I was wondering if there is a better way.