Merge branch 'main' into correct-definition-torchmetrics

ourownstory · Aug 9, 2023 · b1c50ee · b1c50ee
2 parents 5323f23 + 90f6e6e
commit b1c50ee
Show file tree

Hide file tree

Showing 3 changed files with 191 additions and 177 deletions.
diff --git a/neuralprophet/plot_model_parameters_matplotlib.py b/neuralprophet/plot_model_parameters_matplotlib.py
@@ -265,9 +265,12 @@ def plot_trend(m, quantile, ax=None, plot_name="Trend", figsize=(10, 6), df_name
             trend_1 = trend_0
         else:
             if m.model.config_trend.trend_global_local == "local":
-                trend_1 = trend_0 + m.model.trend.trend_k0[quantile_index, m.model.id_dict[df_name]].detach().numpy()
+                trend_1 = (
+                    trend_0
+                    + m.model.trend.trend_k0[quantile_index, m.model.id_dict[df_name]].detach().numpy().squeeze()
+                )
             else:
-                trend_1 = trend_0 + m.model.trend.trend_k0[quantile_index, 0].detach().numpy()
+                trend_1 = trend_0 + m.model.trend.trend_k0[quantile_index, 0].detach().numpy().squeeze()
 
         data_params = m.config_normalization.get_data_params(df_name)
         shift = data_params["y"].shift