I get a runtime error from my dataloader, which loads MNIST, when i use more than one class. When I use only one class, it works.
This is my dataset class (meant for self-supervised learning):
class mnistDataset(Dataset):
    def __init__(self, mnist_data, numbers_to_include=None):

        # Subset the digits to include
        if numbers_to_include is not None:
            data_temp = [im for im in mnist_data if im[1] in numbers_to_include]            
            im_true_temp = [sample[0] for sample in data_temp]
            self.target = [sample[1] for sample in data_temp]
        else:
            im_true_temp = [sample[0] for sample in mnist_data]
            self.target = [sample[1] for sample in mnist_data]
        
        # Remove pixels in the middle and cast to tensors
        self.im = []
        self.im_true = []
        for image in im_true_temp:
            image_temp = image.clone()
            image_temp[0,12:15,:] = 0                
            self.im.append(image_temp)
            self.im_true.append(image)
            
    def __len__(self):
        return len(self.target)
    
    def __getitem__(self, idx):
        return self.im[idx], self.im_true[idx], torch.Tensor(self.target[idx]).type(torch.LongTensor)

I wrap this in the Dataloader class:
batch_size = 16
numbers_to_include = [0,5]

# Dataset 
train_dataset = mnistDataset(trainset, numbers_to_include=numbers_to_include)
train_dataloader = DataLoader(train_dataset, shuffle=True, batch_size=batch_size, drop_last=False)

The error happens when i try to iterate the dataloader (I can iterate the dataset class):
for X, y, t in test_dataloader:
...

This gives the following error:
Traceback (most recent call last):

  File "C:\Users\Malte\Documents\CMI\python\networks\self_supervised_mnist.py", line 112, in <module>
    for X, y, t in test_dataloader:

  File "C:\Users\Malte\Anaconda3\lib\site-packages\torch\utils\data\dataloader.py", line 345, in __next__
    data = self._next_data()

  File "C:\Users\Malte\Anaconda3\lib\site-packages\torch\utils\data\dataloader.py", line 385, in _next_data
    data = self._dataset_fetcher.fetch(index)  # may raise StopIteration

  File "C:\Users\Malte\Anaconda3\lib\site-packages\torch\utils\data\_utils\fetch.py", line 47, in fetch
    return self.collate_fn(data)

  File "C:\Users\Malte\Anaconda3\lib\site-packages\torch\utils\data\_utils\collate.py", line 79, in default_collate
    return [default_collate(samples) for samples in transposed]

  File "C:\Users\Malte\Anaconda3\lib\site-packages\torch\utils\data\_utils\collate.py", line 79, in <listcomp>
    return [default_collate(samples) for samples in transposed]

  File "C:\Users\Malte\Anaconda3\lib\site-packages\torch\utils\data\_utils\collate.py", line 55, in default_collate
    return torch.stack(batch, 0, out=out)

RuntimeError: invalid argument 0: Sizes of tensors must match except in dimension 0. Got 0 and 5 in dimension 1 at C:\w\1\s\tmp_conda_3.7_100118\conda\conda-bld\pytorch_1579082551706\work\aten\src\TH/generic/THTensor.cpp:612

Does anyone have an idea what causes this?