Add sigmoid

luxonis · Nov 18, 2024 · 8216fce · 8216fce
1 parent 13663ed
commit 8216fce
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/luxonis_train/nodes/heads/ddrnet_segmentation_head.py b/luxonis_train/nodes/heads/ddrnet_segmentation_head.py
@@ -105,6 +105,6 @@ def forward(self, inputs: Tensor) -> Tensor:
         x = self.conv2(x)
         x = self.upscale(x)
         if self.export:
-            x = x.argmax(dim=1) if self.n_classes > 1 else (x > 0)
+            x = x.argmax(dim=1) if self.n_classes > 1 else torch.sigmoid(x) > 0.5
             return x.to(dtype=torch.int32)
         return x