Fix 361

autoPyTorch/pipeline/components/training/trainer/__init__.py

@@ -293,6 +293,13 @@ def _fit(self, X: Dict[str, Any], y: Any = None, **kwargs: Any) -> 'TrainerChoic
                 writer=writer,
             )
 
+            # its fine if train_loss is None due to `is_max_time_reached()`
+            if train_loss is None:
+                if self.budget_tracker.is_max_time_reached():
+                    break
+                else:
+                    raise RuntimeError("`train_loss` computed to None.")
+
             val_loss, val_metrics, test_loss, test_metrics = None, {}, None, {}
             if self.eval_valid_each_epoch(X):
                 val_loss, val_metrics = self.choice.evaluate(X['val_data_loader'], epoch, writer)

autoPyTorch/pipeline/components/training/trainer/base_trainer.py

@@ -277,7 +277,7 @@ def _scheduler_step(
 
     def train_epoch(self, train_loader: torch.utils.data.DataLoader, epoch: int,
                     writer: Optional[SummaryWriter],
-                    ) -> Tuple[float, Dict[str, float]]:
+                    ) -> Tuple[Optional[float], Dict[str, float]]:
         """
         Train the model for a single epoch.
 
@@ -317,6 +317,9 @@ def train_epoch(self, train_loader: torch.utils.data.DataLoader, epoch: int,
                     epoch * len(train_loader) + step,
                 )
 
+        if N == 0:
+            return None, {}
+
         self._scheduler_step(step_interval=StepIntervalUnit.epoch, loss=loss_sum / N)
 
         if self.metrics_during_training: