Updated load_data.py to fit with one MNIST dataloader, but to use both MNIST cases

Author: Johanmkr

utils/load_data.py

@@ -3,7 +3,7 @@
 
 from .dataloaders import (
     Downloader,
-    MNISTDataset0_3,
+    MNISTDataset,
     SVHNDataset,
     USPSDataset0_6,
     USPSH5_Digit_7_9_Dataset,
@@ -57,15 +57,17 @@ def load_data(dataset: str, *args, **kwargs) -> tuple:
             train_labels, test_labels = downloader.usps(data_dir=data_dir)
             labels = np.arange(7, 10)
         case "mnist_0-3":
-            dataset = MNISTDataset0_3
+            dataset = MNISTDataset
             train_labels, test_labels = downloader.mnist(data_dir=data_dir)
             labels = np.arange(4)
         case "svhn":
             dataset = SVHNDataset
             train_labels, test_labels = downloader.svhn(data_dir=data_dir)
             labels = np.arange(10)
         case "mnist_4-9":
-            raise NotImplementedError("MNIST 4-9 dataset not yet implemented.")
+            dataset = MNISTDataset
+            train_labels, test_labels = downloader.mnist(data_dir=data_dir)
+            labels = np.arange(4,10)
         case _:
             raise NotImplementedError(f"Dataset: {dataset} not implemented.")