From 302b00b5f4b93bb6cdb3c651dc9f06b66d06016d Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Fri, 8 Apr 2022 12:55:16 +0200 Subject: [PATCH] Update `_make_grid()` (#7346) --- models/yolo.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/models/yolo.py b/models/yolo.py index 85c5a969..f072aeeb 100644 --- a/models/yolo.py +++ b/models/yolo.py @@ -77,13 +77,15 @@ def forward(self, x): def _make_grid(self, nx=20, ny=20, i=0): d = self.anchors[i].device + t = self.anchors[i].dtype shape = 1, self.na, ny, nx, 2 # grid shape + y, x = torch.arange(ny, device=d, dtype=t), torch.arange(nx, device=d, dtype=t) if check_version(torch.__version__, '1.10.0'): # torch>=1.10.0 meshgrid workaround for torch>=0.7 compatibility - yv, xv = torch.meshgrid(torch.arange(ny, device=d), torch.arange(nx, device=d), indexing='ij') + yv, xv = torch.meshgrid(y, x, indexing='ij') else: - yv, xv = torch.meshgrid(torch.arange(ny, device=d), torch.arange(nx, device=d)) - grid = torch.stack((xv, yv), 2).expand(shape).float() - 0.5 # add grid offset, i.e. y = 2.0 * x - 0.5 - anchor_grid = (self.anchors[i] * self.stride[i]).view((1, self.na, 1, 1, 2)).expand(shape).float() + yv, xv = torch.meshgrid(y, x) + grid = torch.stack((xv, yv), 2).expand(shape) - 0.5 # add grid offset, i.e. y = 2.0 * x - 0.5 + anchor_grid = (self.anchors[i] * self.stride[i]).view((1, self.na, 1, 1, 2)).expand(shape) return grid, anchor_grid