xclele/yolov10
1
0
Fork 0
mirror of https://github.com/THU-MIG/yolov10.git synced 2025-05-23 05:24:22 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Small fix in detection head for enabling ov export from scratch (#183)

Browse Source
This commit is contained in:
Ekaterina Aidova 2024-06-03 13:59:44 +04:00 committed by GitHub
parent 1cfe7a4e13
commit 060af8b8c1
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
ultralytics/nn/modules/head.py
View File

@ -520,7 +520,7 @@ class v10Detect(Detect):
else:
assert(self.max_det != -1)
boxes, scores, labels = ops.v10postprocess(one2one.permute(0, 2, 1), self.max_det, self.nc)
return torch.cat([boxes, scores.unsqueeze(-1), labels.unsqueeze(-1)], dim=-1)
return torch.cat([boxes, scores.unsqueeze(-1), labels.unsqueeze(-1).to(boxes.dtype)], dim=-1)
else:
return {"one2many": one2many, "one2one": one2one}