Update nacl_loss.py

Call contiguous after permute to avoid reshaping issue.
Bala93 · Aug 15, 2024 · f579763 · f579763
1 parent db9daeb
commit f579763
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/monai/losses/nacl_loss.py b/monai/losses/nacl_loss.py
@@ -95,11 +95,11 @@ def get_constr_target(self, mask: torch.Tensor) -> torch.Tensor:
         rmask: torch.Tensor
 
         if self.dim == 2:
-            oh_labels = F.one_hot(mask.to(torch.int64), num_classes=self.nc).contiguous().permute(0, 3, 1, 2).float()
+            oh_labels = F.one_hot(mask.to(torch.int64), num_classes=self.nc).permute(0, 3, 1, 2).contiguous().float()
             rmask = self.svls_layer(oh_labels)
 
         if self.dim == 3:
-            oh_labels = F.one_hot(mask.to(torch.int64), num_classes=self.nc).contiguous().permute(0, 4, 1, 2, 3).float()
+            oh_labels = F.one_hot(mask.to(torch.int64), num_classes=self.nc).permute(0, 4, 1, 2, 3).contiguous().float()
             rmask = self.svls_layer(oh_labels)
 
         return rmask