fix flowstate h=1 crash

tests/models/foundation/test_flowstate.py

@@ -0,0 +1,22 @@
+import numpy as np
+import pandas as pd
+
+from timecopilot import TimeCopilotForecaster
+from timecopilot.models.foundation.flowstate import FlowState
+
+
+def test_flowstate_h1_single_uid():
+    # create simple weekly data for one unique_id
+    ds = pd.date_range("2024-01-01", periods=20, freq="W")
+    df = pd.DataFrame({"unique_id": "u1", "ds": ds, "y": np.arange(20)})
+
+    tcf = TimeCopilotForecaster(models=[FlowState()])
+
+    # this used to crash before the fix
+    fcst = tcf.forecast(df=df, h=1, freq="W")
+
+    # basic checks
+    assert isinstance(fcst, pd.DataFrame)
+    assert len(fcst) == 1
+    assert "unique_id" in fcst.columns
+    assert "ds" in fcst.columns

timecopilot/models/foundation/flowstate.py

@@ -167,8 +167,8 @@ def _predict_batch(
             batch_first=False,
         ).prediction_outputs
         fcst = fcst.squeeze(-1).transpose(-1, -2)  # now shape is (batch, h, quantiles)
-        fcst_mean = fcst[..., supported_quantiles.index(0.5)].squeeze()
-        fcst_mean_np = fcst_mean.detach().numpy()
+        fcst_mean = fcst[..., supported_quantiles.index(0.5)]
+        fcst_mean_np = fcst_mean.detach().numpy(force=True)
         fcst_quantiles_np = fcst.detach().numpy() if quantiles is not None else None
         return fcst_mean_np, fcst_quantiles_np