Fix

tinkoff-ai · Oct 8, 2021 · 6d894fc · 6d894fc
1 parent c963ad9
commit 6d894fc
Showing 1 changed file with 1 addition and 3 deletions.
diff --git a/etna/ensembles/voting_ensemble.py b/etna/ensembles/voting_ensemble.py
@@ -124,9 +124,7 @@ def _forecast_pipeline(pipeline: Pipeline) -> TSDataset:
 
     def _vote(self, forecasts: List[TSDataset]) -> TSDataset:
         """Get average forecast."""
-        forecast_df = forecasts[0][:, :, "target"] * self.weights[0]
-        for forecast, weight in zip(forecasts[1:], self.weights[1:]):
-            forecast_df += forecast[:, :, "target"] * weight
+        forecast_df = sum([forecast[:, :, "target"] * weight for forecast, weight in zip(forecasts, self.weights)])
         forecast_dataset = TSDataset(df=forecast_df, freq=forecasts[0].freq)
         return forecast_dataset