[Python] Discrepancy in number of elements outputted by torch Dataset and DataLoader

I have a custom Subset:

class TestSubset2(Subset):
def __init__(self, dataset, indices, days=False):
super().__init__(dataset, indices)
self.days = days
def __getitem__(self, idx):
original_idx = self.indices[idx]
coordinates = self.dataset.all_indices[original_idx]
day = coordinates[0]

if self.dataset.transform:
original_transform = self.dataset.transform
self.dataset.transform = None
image, label = self.dataset[original_idx]
self.dataset.transform = original_transform
else:
image, label = self.dataset[original_idx]

if self.days:
result = (image, label, torch.tensor(day, dtype=torch.int64))
assert len(result) == 3, f"Got {len(result)} pieces at idx={idx}"
return result
else:
result = (image, label)
return result


Here, my Subset/Dataset is clearly set on returning three items provided self.days = True, and this is confirmed by my later usage:

test_dataset = TestSubset2(data, test_idx, days=True), where assert all(len(test_dataset) == 3 for i in range(len(test_dataset))) executes without any problems. Yet when it comes to using my dataloader: test_loader = DataLoader(test_dataset, batch_size=1, shuffle=False), print(len(next(iter(test_loader)))) , the output of the cell is 2, not 3. It seems that the dataloader is not bothering to return the last item.

Continue reading...