Merge pull request selimsef#6 from vessemer/master

BCEDiceJaccardLoss fixes
ljyljyok · Sep 1, 2018 · cc21866 · cc21866
2 parents 1b9c320 + 5a94c9c
commit cc21866
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/albu/src/pytorch_utils/loss.py b/albu/src/pytorch_utils/loss.py
@@ -99,10 +99,11 @@ def __init__(self, weights, weight=None, size_average=True):
 
     def forward(self, input, target):
         loss = 0
+        sigmoid_input = torch.sigmoid(input)
         for k, v in self.weights.items():
-            if not v:
+            if not v: 
                 continue
-            val = self.mapping[k](input, target)
+            val = self.mapping[k](input if k == 'bce' else sigmoid_input, target)
             self.values[k] = val
             if k != 'bce':
                 loss += self.weights[k] * (1 - val)