From 060af8b8c1a49deb364c716504a119db15ac56b4 Mon Sep 17 00:00:00 2001 From: Ekaterina Aidova Date: Mon, 3 Jun 2024 13:59:44 +0400 Subject: [PATCH] Small fix in detection head for enabling ov export from scratch (#183) --- 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}