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}