Compare the value of DatasetMode enum, instead of Enum itself directly

src/dataset/__init__.py

@@ -86,7 +86,7 @@ def get_dataset(
     List[BaseNormalsDataset],
 ]:
     if "mixed" == cfg_data_split.name:
-        assert DatasetMode.TRAIN == mode, "Only training mode supports mixed datasets."
+        assert DatasetMode.TRAIN.value == mode.value, "Only training mode supports mixed datasets."
         dataset_ls = [
             get_dataset(_cfg, base_data_dir, mode, **kwargs)
             for _cfg in cfg_data_split.dataset_list

src/dataset/base_depth_dataset.py

@@ -116,7 +116,7 @@ def __len__(self):
 
     def __getitem__(self, index):
         rasters, other = self._get_data_item(index)
-        if DatasetMode.TRAIN == self.mode:
+        if DatasetMode.TRAIN.value == self.mode.value:
             rasters = self._training_preprocess(rasters)
         # merge
         outputs = rasters

src/dataset/base_iid_dataset.py

@@ -99,7 +99,7 @@ def __len__(self):
 
     def __getitem__(self, index):
         rasters, other = self._get_data_item(index)
-        if DatasetMode.TRAIN == self.mode:
+        if DatasetMode.TRAIN.value == self.mode.value:
             rasters = self._training_preprocess(rasters)
         # merge
         outputs = rasters

src/dataset/base_normals_dataset.py

@@ -92,7 +92,7 @@ def __len__(self):
 
     def __getitem__(self, index):
         rasters, other = self._get_data_item(index)
-        if DatasetMode.TRAIN == self.mode:
+        if DatasetMode.TRAIN.value == self.mode.value:
             rasters = self._training_preprocess(rasters)
         # merge
         outputs = rasters