Take a subset of a DataLoader

I want to create a small subset of the IMDB dataset, so I can test if my program as a whole works before running training on the whole dataset. I tried using a Sampler, but this doesn’t seem to work since the dataset is already loaded as an iterator as follows:

train_iter, test_iter = IMDB(root = ".")

Is there a quick way I can do this?

If you are preloading the entire dataset in its __init__ method you would need to change this logic there.
On the other hand, if you are lazily loading each sample in the __getitem__ you can use a Subset or SubsetRandomSampler to only load the desired indices.