Only select first if only one

tingiskhan · Dec 10, 2024 · 212560d · 212560d
1 parent 126a0ab
commit 212560d
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/skyro/sktime.py b/skyro/sktime.py
@@ -167,6 +167,9 @@ def _predict_proba(self, fh, X, marginal=True):
         if predictions.name is not None:
             warnings.warn(f"Name of the frame will be overwritten with '{columns}'!")
 
+        if len(columns) == 1:
+            columns = columns[0]
+
         as_frame = predictions.to_dataframe(columns)
 
         if predictions.ndim > 2: