[Fix] Problem of one hot when CELoss has weight (#2050)

PaddlePaddle · May 10, 2022 · 6fa1eb0 · 6fa1eb0
1 parent 4fec7b3
commit 6fa1eb0
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/paddleseg/models/losses/cross_entropy_loss.py b/paddleseg/models/losses/cross_entropy_loss.py
@@ -119,7 +119,7 @@ def _post_process_loss(self, logit, label, semantic_weights, loss):
             loss = loss * semantic_weights
 
         if self.weight is not None:
-            _one_hot = F.one_hot(label, logit.shape[-1])
+            _one_hot = F.one_hot(label * mask, logit.shape[-1])
             coef = paddle.sum(_one_hot * self.weight, axis=-1)
         else:
             coef = paddle.ones_like(label)