chore: add assert statement for dataset name

AFAgarap · Nov 2, 2020 · 7bcf0fe · 7bcf0fe
1 parent 2ea0582
commit 7bcf0fe
Showing 1 changed file with 8 additions and 1 deletion.
diff --git a/pt_datasets/load_dataset.py b/pt_datasets/load_dataset.py
@@ -45,9 +45,16 @@ def load_dataset(
     Tuple[object, object]
         A tuple consisting of the training dataset and the test dataset.
     """
-    transform = torchvision.transforms.Compose([torchvision.transforms.ToTensor()])
+    supported_datasets = ["mnist", "fashion_mnist", "emnist", "cifar10", "svhn"]
+
     name = name.lower()
 
+    assert (
+        name in supported_datasets
+    ), f"[ERROR] Dataset {name} is not supported. Supported datasets: mnist, fashion_mnist, emnist, cifar10, svhn."
+
+    transform = torchvision.transforms.Compose([torchvision.transforms.ToTensor()])
+
     if name == "mnist":
         train_dataset = torchvision.datasets.MNIST(
             root=data_folder, train=True, download=True, transform=transform