[Fix] Fix the error of BCE loss when batch size is 1. (open-mmlab#1629)

pytorchuser · Apr 11, 2023 · 62f44c0 · 62f44c0
1 parent 1c58d05
commit 62f44c0
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/mmseg/models/losses/cross_entropy_loss.py b/mmseg/models/losses/cross_entropy_loss.py
@@ -124,7 +124,7 @@ def binary_cross_entropy(pred,
         assert label[label != ignore_index].max() <= 1, \
             'For pred with shape [N, 1, H, W], its label must have at ' \
             'most 2 classes'
-        pred = pred.squeeze()
+        pred = pred.squeeze(1)
     if pred.dim() != label.dim():
         assert (pred.dim() == 2 and label.dim() == 1) or (
                 pred.dim() == 4 and label.dim() == 3), \