address comment

awslabs · Jan 29, 2023 · e8bfc9a · e8bfc9a
1 parent 701328f
commit e8bfc9a
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/src/gluonts/mx/model/n_beats/_ensemble.py b/src/gluonts/mx/model/n_beats/_ensemble.py
@@ -189,9 +189,9 @@ def predict(
             # default according to paper is median,
             # but we can also make use of not aggregating
             if self.aggregation_method == "median":
-                output = np.median(output, axis=0)
+                output = np.median(output, axis=0, keepdims=True)
             elif self.aggregation_method == "mean":
-                output = np.mean(output, axis=0)
+                output = np.mean(output, axis=0, keepdims=True)
             else:  # "none": do not aggregate
                 pass