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

Small fix in detection head for enabling ov export from scratch

Browse Source
This commit is contained in:
eaidova 2024-06-03 13:35:41 +04:00
parent 1cfe7a4e13
commit 1d01fcb096
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: else:
assert(self.max_det != -1) assert(self.max_det != -1)
boxes, scores, labels = ops.v10postprocess(one2one.permute(0, 2, 1), self.max_det, self.nc) 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: else:
return {"one2many": one2many, "one2one": one2one} return {"one2many": one2many, "one2one": one2one}