Fix export files

2023-09-12 15:44:13 -03:00
parent b2388fb3cb
commit b738644f6e
7 changed files with 7 additions and 7 deletions
--- a/utils/export_yoloV6.py
+++ b/utils/export_yoloV6.py
@@ -22,7 +22,7 @@ class DeepStreamOutput(nn.Module):
    def forward(self, x):
        boxes = x[:, :, :4]
        objectness = x[:, :, 4:5]
-        scores, classes = torch.max(x[:, :, 5:6], 2, keepdim=True)
+        scores, classes = torch.max(x[:, :, 5:], 2, keepdim=True)
        scores *= objectness
        classes = classes.float()
        return boxes, scores, classes