Merge branch 'ecmwf:develop' into pr/aw_rescale

metno · Nov 26, 2024 · 7bb2919 · 7bb2919
2 parents c159e27 + 1abb65e
commit 7bb2919
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 1 deletion.
diff --git a/src/anemoi/training/losses/weightedloss.py b/src/anemoi/training/losses/weightedloss.py
@@ -107,6 +107,8 @@ def scale(
 
         if scalar_indices is None:
             return x * scalar
+
+        scalar = scalar.expand_as(x)
         return x * scalar[scalar_indices]
 
     def scale_by_node_weights(self, x: torch.Tensor, squash: bool = True) -> torch.Tensor:

diff --git a/src/anemoi/training/train/forecaster.py b/src/anemoi/training/train/forecaster.py
@@ -553,7 +553,7 @@ def calculate_val_metrics(
                 metrics[f"{metric_name}/{mkey}/{rollout_step + 1}"] = metric(
                     y_pred_postprocessed[..., indices],
                     y_postprocessed[..., indices],
-                    scalar_indices=[..., indices],
+                    scalar_indices=[..., indices] if -1 in metric.scalar else None,
                 )
 
         return metrics