Major refactor, small fixes

data_aug/contrastive_learning_dataset.py

@@ -2,6 +2,7 @@
 from data_aug.gaussian_blur import GaussianBlur
 from torchvision import transforms, datasets
 from data_aug.view_generator import ContrastiveLearningViewGenerator
+from exceptions.exceptions import InvalidDatasetSelection
 
 
 class ContrastiveLearningDataset:
@@ -33,5 +34,9 @@ def get_dataset(self, name, n_views):
                                                               n_views),
                                                           download=True)}
 
-        dataset = valid_datasets.get(name, 'Invalid dataset option.')()
-        return dataset
+        try:
+            dataset_fn = valid_datasets[name]
+        except KeyError:
+            raise InvalidDatasetSelection()
+        else:
+            return dataset_fn()

exceptions/exceptions.py

@@ -3,4 +3,8 @@ class BaseSimCLRException(Exception):
 
 
 class InvalidBackboneError(BaseSimCLRException):
-    """Raised when the choice of backbone Convnet is invalid."""
+    """Raised when the choice of backbone Convnet is invalid."""
+
+
+class InvalidDatasetSelection(BaseSimCLRException):
+    """Raised when the choice of dataset is invalid."""