Return self in _fit()

Author: franchuterivera

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

@@ -203,7 +203,7 @@ def fit(self, X: Dict[str, Any], y: Any = None, **kwargs: Any) -> autoPyTorchCom
 
         return cast(autoPyTorchComponent, self.choice)
 
-    def _fit(self, X: Dict[str, Any], y: Any = None, **kwargs: Any) -> torch.nn.Module:
+    def _fit(self, X: Dict[str, Any], y: Any = None, **kwargs: Any) -> 'TrainerChoice':
         """
         Fits a component by using an input dictionary with pre-requisites
 
@@ -336,7 +336,7 @@ def _fit(self, X: Dict[str, Any], y: Any = None, **kwargs: Any) -> torch.nn.Modu
         # Tag as fitted
         self.fitted_ = True
 
-        return X['network'].state_dict()
+        return self
 
     def early_stop_handler(self, X: Dict[str, Any]) -> bool:
         """