adapt TFT tests

unit8co · May 4, 2022 · f6711f5 · f6711f5
1 parent 0eeba80
commit f6711f5
Showing 1 changed file with 7 additions and 3 deletions.
diff --git a/darts/tests/models/forecasting/test_TFT.py b/darts/tests/models/forecasting/test_TFT.py
@@ -282,10 +282,14 @@ def helper_fit_predict(
                 series=series,
                 past_covariates=past_covariates,
                 future_covariates=future_covariates,
-                num_samples=100,
+                num_samples=(100 if model._is_probabilistic() else 1),
             )
+
             if isinstance(y_hat, TimeSeries):
-                y_hat = y_hat.quantile_timeseries(0.5)
+                y_hat = y_hat.quantile_timeseries(0.5) if y_hat.n_samples > 1 else y_hat
             else:
-                y_hat = [ts.quantile_timeseries(0.5) for ts in y_hat]
+                y_hat = [
+                    ts.quantile_timeseries(0.5) if ts.n_samples > 1 else ts
+                    for ts in y_hat
+                ]
             return y_hat