Fix bug in to_quantile_forecast (#2225)

awslabs · Aug 26, 2022 · 259481f · 259481f
1 parent 48231da
commit 259481f
Showing 1 changed file with 7 additions and 4 deletions.
diff --git a/src/gluonts/model/forecast.py b/src/gluonts/model/forecast.py
@@ -636,11 +636,14 @@ def __repr__(self):
  ]
  )
 
- def to_quantile_forecast(
- self, quantiles: List[Union[float, str]]
- ) -> "QuantileForecast":
+ def to_quantile_forecast(self, quantiles: List[str]) -> "QuantileForecast":
  return QuantileForecast(
- forecast_arrays=np.array([self.quantile(q) for q in quantiles]),
+ forecast_arrays=np.array(
+ [
+ self.quantile(q) if q != "mean" else self.mean()
+ for q in quantiles
+ ]
+ ),
  start_date=self.start_date,
  forecast_keys=quantiles,
  item_id=self.item_id,