From a2336088f0dbedc421774660534fe10d7a3da840 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Thu, 4 Jun 2020 14:25:10 -0700 Subject: [PATCH] updates --- models/yolo.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/models/yolo.py b/models/yolo.py index 7e36dab..0457db6 100644 --- a/models/yolo.py +++ b/models/yolo.py @@ -27,7 +27,7 @@ class Detect(nn.Module): x[i] = x[i].view(bs, self.na, self.no, ny, nx).permute(0, 1, 3, 4, 2).contiguous() if not self.training: # inference - if self.grid[i].shape[2:4] != x[i].shape[2:4]: + if (self.grid[i].shape[2:4] != x[i].shape[2:4]) | (self.grid[i].device != x[i].device): self.grid[i] = self._make_grid(nx, ny).to(x[i].device) y = x[i].sigmoid()