Change col names

sktime · Nov 22, 2024 · 270be35 · 270be35
1 parent d765033
commit 270be35
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/skpro/regression/bayesian/bayesian_conjugate.py b/skpro/regression/bayesian/bayesian_conjugate.py
@@ -94,6 +94,8 @@ def _fit(self, X, y):
         if not is_centered:
             X -= X.mean(axis=0)
 
+        self._y_cols = y.columns
+
         # Construct the prior mean, covariance, and precision
         self._prior_mu = np.zeros(n_features)
         self._prior_cov = np.eye(n_features) / self.alpha
@@ -116,7 +118,6 @@ def _predict_proba(self, X):
         y_pred : Normal
             Predicted Normal distribution for outputs.
         """
-        y_cols = self._y_cols
         idx = X.index
         if isinstance(X, pd.DataFrame):
             X = X.values
@@ -145,7 +146,7 @@ def _predict_proba(self, X):
         self._y_pred = Normal(
             mu=mus,
             sigma=sigmas,
-            columns=y_cols,
+            columns=self.y_cols,
             index=idx,
         ).iloc[:, 0]
         return self._y_pred