fix bce loss

janelia-cellmap · Sep 27, 2023 · cd4077d · cd4077d
1 parent ee03505
commit cd4077d
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/dacapo/experiments/tasks/losses/hot_distance_loss.py b/dacapo/experiments/tasks/losses/hot_distance_loss.py
@@ -13,8 +13,9 @@ def compute(self, prediction, target, weight):
         return self.hot_loss(prediction_hot, target_hot, weight_hot) + self.distance_loss(prediction_distance, target_distance, weight_distance)
 
     def hot_loss(self, prediction, target, weight):
-        loss = torch.nn.BCELoss()
-        return loss(prediction * weight, target * weight)
+        loss = torch.nn.BCEWithLogitsLoss(reduction='none')
+        return torch.mean(loss(prediction , target) * weight)
+        # return abs(prediction * weight - target * weight).sum()
 
     def distance_loss(self, prediction, target, weight):
         loss = torch.nn.MSELoss()