Upload service.py

service.py

@@ -95,7 +95,11 @@ class ChronosService:
         assert self._pipeline is not None
 
         torch = self._torch
-        context = torch.tensor(close_prices[-self.max_context_length :], dtype=torch.float32)
+        # Chronos-2 expects (n_series, n_variates, history_length) for tensor input.
+        context = torch.tensor(
+            close_prices[-self.max_context_length :],
+            dtype=torch.float32,
+        ).reshape(1, 1, -1)
         forecast = self._pipeline.predict(
             context,
             prediction_length=self.max_horizon_step,
@@ -127,7 +131,11 @@ class ChronosService:
         else:
             tensor = torch.as_tensor(forecast)
 
-        if tensor.ndim == 3:
+        if tensor.ndim == 4:
+            samples = tensor[0, :, 0, :]
+            mean_forecast = samples.mean(dim=0).tolist()
+            std_forecast = samples.std(dim=0).tolist()
+        elif tensor.ndim == 3:
             samples = tensor[0]
             mean_forecast = samples.mean(dim=0).tolist()
             std_forecast = samples.std(dim=0).tolist()