Torch Dataset Looping too far
Why does this dataset try to iterate past the final element

from torch.utils.data.dataset import Dataset
class DumbDataset(Dataset):
    def __init__(self, dct):
        self.dct = dct
        self.mapping = dict(enumerate(dct))
    def __getitem__(self, index):
        return self.dct[self.mapping[index]]

    def __len__(self):
        print('called')
        return len(self.dct)

ds = DumbDataset({'a': 'aword', 'b': 'another_words'})

for k in ds: print(k)


This raises KeyError: 2, which I don't understand since length of the object is 2. Shouldn't the iterator get StopIteration once it is exhausted?

The reason why your code raises KeyError is that Dataset does not implement __iter__() and thus when used in a for-loop Python falls back to starting at index 0 and calling __getitem__ until IndexError is raised, as discussed here. You can modify DumbDataset to work like this by having it raise an IndexError when the index is out of bounds

def __getitem__(self, index):
    if index &gt;= len(self): raise IndexError
    return self.dct[self.mapping[index]]


and then your loop

for k in ds:
    print(k)


will work as you expected. On the other hand, the typical template for torch datasets is that you can either loop through them with indexing

for i in range(len(ds)):
    k = ds[k]
    print(k)


or that you wrap them in a DataLoader which returns elements in batches

generator = DataLoader(ds)
for k in generator:
    print(k)

