From 618d1e195e19d13086da322cb7d8f313741d468d Mon Sep 17 00:00:00 2001 From: SeeFun Date: Thu, 5 Aug 2021 16:03:29 +0800 Subject: [PATCH] fix bug in lovasz loss https://github.com/BloodAxe/pytorch-toolbelt/issues/61 --- pytorch_toolbelt/losses/lovasz.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pytorch_toolbelt/losses/lovasz.py b/pytorch_toolbelt/losses/lovasz.py index c83cf874e..91a938c4d 100644 --- 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)