Change DDRNet Head Output Dtype (#88)

luxonis · Oct 3, 2024 · 6a6f2db · 6a6f2db
1 parent ca9ab66
commit 6a6f2db
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/luxonis_train/nodes/heads/ddrnet_segmentation_head.py b/luxonis_train/nodes/heads/ddrnet_segmentation_head.py
@@ -83,13 +83,13 @@ def __init__(
         )
 
     def forward(self, inputs: Tensor) -> Tensor:
-        x = self.relu(self.bn1(inputs))
+        x: Tensor = self.relu(self.bn1(inputs))
         x = self.conv1(x)
         x = self.relu(self.bn2(x))
         x = self.conv2(x)
         x = self.upscale(x)
         if self.export:
-            return x.argmax(dim=1)
+            return x.argmax(dim=1).to(dtype=torch.int32)
         return x
 
     def set_export_mode(self, mode: bool = True) -> None: