fix flake

Author: ravinkohli

autoPyTorch/api/base_task.py

@@ -1491,6 +1491,10 @@ def fit_pipeline(
                                        dataset_name=dataset_name
                                        )
 
+        # dataset_name is created inside the constructor of BaseDataset
+        # we expect it to be not None. This is for mypy
+        assert dataset.dataset_name is not None
+
         # TAE expects each configuration to have a config_id.
         # For fitting a pipeline as it is not part of the
         # search process, it makes sense to set it to 0
@@ -1506,9 +1510,6 @@ def fit_pipeline(
         self._backend.save_datamanager(dataset)
 
         if self._logger is None:
-            # dataset_name is created inside the constructor of BaseDataset
-            # we expect it to be not None. This is for mypy
-            assert dataset.dataset_name is not None
             self._logger = self._get_logger(dataset.dataset_name)
 
         include_components = self.include_components if include_components is None else include_components
@@ -1576,6 +1577,7 @@ def fit_pipeline(
         )
 
         fitted_pipeline = self._get_fitted_pipeline(
+            dataset_name=dataset.dataset_name,
             pipeline_idx=run_info.config.config_id + tae.initial_num_run,
             run_info=run_info,
             run_value=run_value,
@@ -1588,11 +1590,16 @@ def fit_pipeline(
 
     def _get_fitted_pipeline(
         self,
+        dataset_name: str,
         pipeline_idx: int,
         run_info: RunInfo,
         run_value: RunValue,
         disable_file_output: List[Union[str, DisableFileOutputParameters]]
     ) -> Optional[BasePipeline]:
+
+        if self._logger is None:
+            self._logger = self._get_logger(str(dataset_name))
+
         if run_value.status != StatusType.SUCCESS:
             warnings.warn(f"Fitting pipeline failed with status: {run_value.status}"
                           f", additional_info: {run_value.additional_info}")
@@ -1606,7 +1613,7 @@ def _get_fitted_pipeline(
         else:
             load_function = self._backend.load_model_by_seed_and_id_and_budget
 
-        return load_function(
+        return load_function(  # type: ignore[no-any-return]
             seed=self.seed,
             idx=pipeline_idx,
             budget=float(run_info.budget),

autoPyTorch/evaluation/abstract_evaluator.py

@@ -858,7 +858,7 @@ def file_output(
                 self.backend.save_targets_ensemble(self.Y_optimization)
 
         if getattr(self, 'pipelines', None) is not None:
-            if self.pipelines[0] is not None and len(self.pipelines) > 0:
+            if self.pipelines[0] is not None and len(self.pipelines) > 0:  # type: ignore[index, arg-type]
                 if 'pipelines' not in self.disable_file_output:
                     if self.task_type in CLASSIFICATION_TASKS:
                         pipelines = VotingClassifier(estimators=None, voting='soft', )

test/test_utils/test_common.py

@@ -1,6 +1,8 @@
 """
 This tests the functionality in autoPyTorch/utils/common.
 """
+from enum import Enum
+
 import pytest
 
 from autoPyTorch.utils.common import autoPyTorchEnum
@@ -11,6 +13,10 @@ class SubEnum(autoPyTorchEnum):
     y = "y"
 
 
+class DummyEnum(Enum):  # You need to move it on top
+    x = "x"
+
+
 @pytest.mark.parametrize('iter',
                          ([SubEnum.x],
                           ["x"],
@@ -32,9 +38,6 @@ def test_autopytorch_enum(iter):
 
     assert e in iter
 
-class DummyEnum(Enum):  # You need to move it on top
-    x = "x"
-
 
 @pytest.mark.parametrize('iter',
                          [[SubEnum.y],
@@ -67,5 +70,3 @@ def test_raise_errors_autopytorch_enum(others):
 
     with pytest.raises(RuntimeError):
         SubEnum.x == others
-
-