Fixed BCE loss (hopefully)..

metno · Dec 18, 2024 · d13af3d · d13af3d
1 parent c22cbaa
commit d13af3d
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/src/anemoi/training/losses/binarycrossentropy.py b/src/anemoi/training/losses/binarycrossentropy.py
@@ -14,6 +14,8 @@
 
 import torch
 
+from torch.nn import BCELoss
+
 from anemoi.training.losses.weightedloss import BaseWeightedLoss
 
 LOGGER = logging.getLogger(__name__)
@@ -76,6 +78,5 @@ def forward(
         torch.Tensor
             Weighted binary cross entropy loss
         """
-        out = torch.nn.BCELoss(pred, target)
-        out = self.scale(out, scalar_indices, without_scalars=without_scalars)
-        return self.scale_by_node_weights(out, squash)
+        loss = BCELoss()
+        return loss(pred.float(), target.float())  # BCE likes floats.