Merge pull request #62 from seefun/develop

fix bug in lovasz loss
BloodAxe · Aug 5, 2021 · f050f3a · f050f3a
2 parents d022736 + 618d1e1
commit f050f3a
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/pytorch_toolbelt/losses/lovasz.py b/pytorch_toolbelt/losses/lovasz.py
@@ -147,7 +147,7 @@ def _flatten_probas(probas, labels, ignore=None):
         probas = probas.view(B, 1, H, W)
 
     C = probas.size(1)
-    probas = torch.movedim(probas, 0, -1)  # [B, C, Di, Dj, Dk...] -> [B, C, Di...Dk, C]
+    probas = torch.movedim(probas, 1, -1)  # [B, C, Di, Dj, ...] -> [B, Di, Dj, ..., C]
     probas = probas.contiguous().view(-1, C)  # [P, C]
 
     labels = labels.view(-1)