RuntimeError when using more than one class in dataloader

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]
   = [sample[1] for sample in data_temp]
            im_true_temp = [sample[0] for sample in mnist_data]
   = [sample[1] for sample in mnist_data]
        # Remove pixels in the middle and cast to tensors = []
        self.im_true = []
        for image in im_true_temp:
            image_temp = image.clone()
            image_temp[0,12:15,:] = 0                
    def __len__(self):
        return len(
    def __getitem__(self, idx):
        return[idx], self.im_true[idx], torch.Tensor([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\", line 112, in <module>
    for X, y, t in test_dataloader:

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

  File "C:\Users\Malte\Anaconda3\lib\site-packages\torch\utils\data\", 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\", line 47, in fetch
    return self.collate_fn(data)

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

  File "C:\Users\Malte\Anaconda3\lib\site-packages\torch\utils\data\_utils\", 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?