Further code quality improvements

metno · Nov 5, 2024 · ee0d3c9 · ee0d3c9
1 parent 2ee5af4
commit ee0d3c9
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 10 deletions.
diff --git a/src/anemoi/training/losses/mse.py b/src/anemoi/training/losses/mse.py
@@ -136,12 +136,7 @@ def __init__(
         if data_variances is not None:
             self.register_buffer("ivar", data_variances, persistent=True)
 
-    def forward(
-        self,
-        pred: torch.Tensor,
-        target: torch.Tensor,
-        squash=True
-    ) -> torch.Tensor:
+    def forward(self, pred: torch.Tensor, target: torch.Tensor, squash=True) -> torch.Tensor:
         """Calculates the lat-weighted MSE loss.
 
         Parameters

diff --git a/src/anemoi/training/train/forecaster.py b/src/anemoi/training/train/forecaster.py
@@ -110,29 +110,29 @@ def __init__(
             # Define stretched grid metrics according to options for stretched grid loss logging
             self.sg_metrics = torch.nn.ModuleDict()
             if config.diagnostics.sg_metrics.wmse_per_region:
-                self.sg_metrics['wmse_inside_lam_epoch'] = WeightedMSELossStretchedGrid(
+                self.sg_metrics["wmse_inside_lam_epoch"] = WeightedMSELossStretchedGrid(
                     node_weights=self.loss_weights,
                     mask=self.mask,
                     inside_LAM=True,
                     wmse_contribution=False,
                     data_variances=loss_scaling,
                 )
-                self.sg_metrics['wmse_outside_lam_epoch'] = WeightedMSELossStretchedGrid(
+                self.sg_metrics["wmse_outside_lam_epoch"] = WeightedMSELossStretchedGrid(
                     node_weights=self.loss_weights,
                     mask=self.mask,
                     inside_LAM=False,
                     wmse_contribution=False,
                     data_variances=loss_scaling,
                 )
             if config.diagnostics.sg_metrics.wmse_contributions:
-                self.sg_metrics['wmse_inside_lam_contribution_epoch'] = WeightedMSELossStretchedGrid(
+                self.sg_metrics["wmse_inside_lam_contribution_epoch"] = WeightedMSELossStretchedGrid(
                     node_weights=self.loss_weights,
                     mask=self.mask,
                     inside_LAM=True,
                     wmse_contribution=True,
                     data_variances=loss_scaling,
                 )
-                self.sg_metrics['wmse_outside_lam_contribution_epoch'] = WeightedMSELossStretchedGrid(
+                self.sg_metrics["wmse_outside_lam_contribution_epoch"] = WeightedMSELossStretchedGrid(
                     node_weights=self.loss_weights,
                     mask=self.mask,
                     inside_LAM=False,