Update utils.py

Signed-off-by: Dimitris Karkalousos <d.karkalousos@amsterdamumc.nl>
wdika · Nov 3, 2024 · 210bdb4 · 210bdb4
1 parent c8a144b
commit 210bdb4
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/atommic/collections/segmentation/losses/utils.py b/atommic/collections/segmentation/losses/utils.py
@@ -109,4 +109,4 @@ def one_hot(labels: torch.Tensor, num_classes: int, dtype: torch.dtype = torch.f
     sh[dim] = num_classes
     o = torch.zeros(size=sh, dtype=dtype, device=labels.device)
     labels = o.scatter_(dim=dim, index=labels.long(), value=1)
-    return labels
+    return labels