From 1d01fcb096d9f92418965a95e1c4241b47bb8dcf Mon Sep 17 00:00:00 2001 From: eaidova Date: Mon, 3 Jun 2024 13:35:41 +0400 Subject: [PATCH] Small fix in detection head for enabling ov export from scratch --- ultralytics/nn/modules/head.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ultralytics/nn/modules/head.py b/ultralytics/nn/modules/head.py index b55942c2..a9c5d9ee 100644 --- a/ultralytics/nn/modules/head.py +++ b/ultralytics/nn/modules/head.py @@ -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}