Upload service.py

service.py

@@ -131,19 +131,17 @@ class ChronosService:
         else:
             tensor = torch.as_tensor(forecast)
 
-        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()
-        elif tensor.ndim == 2:
-            mean_forecast = tensor[0].tolist()
+        tensor = tensor.to(dtype=torch.float32)
+        squeezed = tensor.squeeze()
+        if squeezed.ndim == 0:
+            raise RuntimeError(f"unexpected Chronos forecast shape: {tuple(tensor.shape)}")
+        if squeezed.ndim == 1:
+            mean_forecast = squeezed.tolist()
             std_forecast = [0.0 for _ in mean_forecast]
         else:
-            raise RuntimeError(f"unexpected Chronos forecast shape: {tuple(tensor.shape)}")
+            samples = squeezed.reshape(-1, squeezed.shape[-1])
+            mean_forecast = samples.mean(dim=0).tolist()
+            std_forecast = samples.std(dim=0).tolist()
 
         confidence: list[float] = []
         for pred, std in zip(mean_forecast, std_forecast):