update cv schema

Nixtla · Nov 22, 2024 · aa016aa · aa016aa
1 parent fb42b70
commit aa016aa
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 4 deletions.
diff --git a/mlforecast/distributed/forecast.py b/mlforecast/distributed/forecast.py
@@ -654,9 +654,8 @@ def cross_validation(
                     keep_last_n=keep_last_n,
                     window_info=window_info,
                 )
-            schema = (
-                self._get_predict_schema()
-                + f",cutoff:datetime,{self._base_ts.target_col}:double"
+            schema = self._get_predict_schema() + Schema(
+                ("cutoff", "datetime"), (self._base_ts.target_col, "double")
             )
             preds = fa.transform(
                 partition_results,

diff --git a/nbs/distributed.forecast.ipynb b/nbs/distributed.forecast.ipynb
@@ -703,7 +703,10 @@
     "                    keep_last_n=keep_last_n,\n",
     "                    window_info=window_info,\n",
     "                )\n",
-    "            schema = self._get_predict_schema() + f',cutoff:datetime,{self._base_ts.target_col}:double'\n",
+    "            schema = (\n",
+    "                self._get_predict_schema() + Schema(\n",
+    "                ('cutoff', 'datetime'), (self._base_ts.target_col, 'double'))\n",
+    "            )\n",
     "            preds = fa.transform(\n",
     "                partition_results,\n",
     "                DistributedMLForecast._predict,\n",