Made one MNIST dataloader since we pass in sample_ids

Author: Johanmkr

utils/dataloaders/__init__.py

@@ -1,13 +1,13 @@
 __all__ = [
     "USPSDataset0_6",
     "USPSH5_Digit_7_9_Dataset",
-    "MNISTDataset0_3",
+    "MNISTDataset",
     "Downloader",
     "SVHNDataset",
 ]
 
 from .download import Downloader
-from .mnist_0_3 import MNISTDataset0_3
+from .mnist import MNISTDataset
 from .svhn import SVHNDataset
 from .usps_0_6 import USPSDataset0_6
 from .uspsh5_7_9 import USPSH5_Digit_7_9_Dataset

utils/dataloaders/mnist_0_3.py utils/dataloaders/mnist.pyutils/dataloaders/mnist_0_3.py renamed to utils/dataloaders/mnist.py

@@ -6,9 +6,9 @@
 from .datasources import MNIST_SOURCE
 
 
-class MNISTDataset0_3(Dataset):
+class MNISTDataset(Dataset):
     """
-    A custom Dataset class for loading a subset of the MNIST dataset containing digits 0 to 3.
+    A custom Dataset class for loading a subset of the MNIST dataset
     Parameters
     ----------
     data_path : Path
@@ -61,7 +61,6 @@ def __init__(
         self.idx = sample_ids
         self.train = train
         self.transform = transform
-        self.num_classes = 4
 
         self.images_path = self.mnist_path / (
             MNIST_SOURCE["train_images"][1] if train else MNIST_SOURCE["test_images"][1]