Fix ONNX export

2023-05-29 21:54:43 -03:00
parent 141c0f2fee
commit b2c4bee8dc
9 changed files with 9 additions and 9 deletions
--- a/utils/export_yoloV7_u6.py
+++ b/utils/export_yoloV7_u6.py
@@ -18,7 +18,7 @@ class DeepStreamOutput(nn.Module):
        x = x.transpose(1, 2)
        boxes = x[:, :, :4]
        scores, classes = torch.max(x[:, :, 4:], 2, keepdim=True)
-        return torch.cat((boxes, scores, classes), dim=2)
+        return torch.cat((boxes, scores, classes.float()), dim=2)


 def suppress_warnings():