override set_output in the classifier and the regressor

Author: LeoGrin

src/tabpfn/classifier.py

@@ -25,6 +25,7 @@
 
 import numpy as np
 import torch
+from sklearn import config_context
 from sklearn.base import BaseEstimator, ClassifierMixin, check_is_fitted
 from sklearn.preprocessing import LabelEncoder
 
@@ -374,6 +375,7 @@ def __sklearn_tags__(self) -> Tags:
         tags.estimator_type = "classifier"
         return tags
 
+    @config_context(transform_output="default")
     def fit(self, X: XType, y: YType) -> Self:
         """Fit the model.
 
@@ -518,6 +520,7 @@ def predict(self, X: XType) -> np.ndarray:
         y = np.argmax(proba, axis=1)
         return self.label_encoder_.inverse_transform(y)  # type: ignore
 
+    @config_context(transform_output="default")
     def predict_proba(self, X: XType) -> np.ndarray:
         """Predict the probabilities of the classes for the provided input samples.
 

src/tabpfn/regressor.py

@@ -25,6 +25,7 @@
 
 import numpy as np
 import torch
+from sklearn import config_context
 from sklearn.base import (
     BaseEstimator,
     RegressorMixin,
@@ -380,6 +381,7 @@ def __sklearn_tags__(self) -> Tags:
         tags.estimator_type = "regressor"
         return tags
 
+    @config_context(transform_output="default")
     def fit(self, X: XType, y: YType) -> Self:
         """Fit the model.
 
@@ -556,6 +558,7 @@ def predict(
     ) -> dict[str, np.ndarray | FullSupportBarDistribution]: ...
 
     # FIXME: improve to not have noqa C901, PLR0912
+    @config_context(transform_output="default")
     def predict(  # noqa: C901, PLR0912
         self,
         X: XType,