I tried to write a custom dataloader for mnist where I want only items with specific labels, and when I try to run my model Cuda gives me out of memory errors after a couple of epochs. When I run my model with the standard MNIST dataloader, the program works fine. Any idea why this is happening?
class MNISTCustomDataset(Dataset):
    def __init__(self, numbers, transform=None, data_dir='./data/'):
        #Training Data
        f = open('./data/train-images-idx3-ubyte')
        loaded = np.fromfile(file=f, dtype=np.uint8)
    
        all_images = loaded[16:].reshape((60000, 28, 28, 1)).astype(np.float32) / 255    
    
    
        f = open('./data/train-labels-idx1-ubyte')
        loaded = np.fromfile(file=f, dtype=np.uint8)
        
        all_labels = loaded[8:].reshape((60000)).astype(np.int32)
        
        self.images = []
        self.labels = []
        for idx in range(0,len(all_images)):
            if all_labels[idx] in numbers:
                self.images.append(all_images[idx])
                self.labels.append(all_labels[idx])
        
        self.transform=transform
    
    def __getitem__(self, index):
        img   = self.images[index]
        label = self.labels[index]

        if self.transform is not None:
            img = self.transform(img)

        return img, label
    
    def __len__(self):
        return len(self.images)

def load_custom_mnist(bsize, numbers):
    dataset = MNISTCustomDataset(numbers, transform=transforms.Compose([
                               transforms.ToTensor(),
                           ]))
    loader  = DataLoader(dataset, batch_size=bsize, shuffle=True)
    
    return loader