From 9043272f78228ce339b60092a1af86c561ae8bdc Mon Sep 17 00:00:00 2001 From: ZiQiangXie <59564630+ZiQiangXie@users.noreply.github.com> Date: Mon, 1 Jan 2024 22:23:28 +0800 Subject: [PATCH] Fix RTDETR generate anchor grid out of boundary (#7247) --- 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 7b04f87f..e7c36ba2 100644 --- a/ultralytics/nn/modules/head.py +++ b/ultralytics/nn/modules/head.py @@ -303,7 +303,7 @@ class RTDETRDecoder(nn.Module): grid_y, grid_x = torch.meshgrid(sy, sx, indexing='ij') if TORCH_1_10 else torch.meshgrid(sy, sx) grid_xy = torch.stack([grid_x, grid_y], -1) # (h, w, 2) - valid_WH = torch.tensor([h, w], dtype=dtype, device=device) + valid_WH = torch.tensor([w, h], dtype=dtype, device=device) grid_xy = (grid_xy.unsqueeze(0) + 0.5) / valid_WH # (1, h, w, 2) wh = torch.ones_like(grid_xy, dtype=dtype, device=device) * grid_size * (2.0 ** i) anchors.append(torch.cat([grid_xy, wh], -1).view(-1, h * w, 4)) # (1, h*w, 4)