fix: requested changes

mindee · Feb 14, 2022 · 7fb3c2b · 7fb3c2b
1 parent 8e6ef60
commit 7fb3c2b
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/doctr/models/detection/linknet/pytorch.py b/doctr/models/detection/linknet/pytorch.py
@@ -184,7 +184,7 @@ def compute_loss(
         dice_loss = 1 - 2 * inter / (cardinality + 1e-8)
 
         # Only consider contributions overlaping the mask
-        return bce_loss[seg_mask].mean() + dice_loss.mean()
+        return bce_loss[seg_mask].mean() + dice_loss
 
 
 def _linknet(

diff --git a/doctr/models/detection/linknet/tensorflow.py b/doctr/models/detection/linknet/tensorflow.py
@@ -164,7 +164,7 @@ def compute_loss(
         cardinality = tf.math.reduce_sum(prob_map[seg_mask] + seg_target[seg_mask])
         dice_loss = 1 - 2 * inter / (cardinality + 1e-8)
 
-        return tf.math.reduce_mean(bce_loss[seg_mask]) + tf.math.reduce_mean(dice_loss)
+        return tf.math.reduce_mean(bce_loss[seg_mask]) + dice_loss
 
     def call(
         self,