zjc
/
yolov5
1
0
Fork 0
Code Issues Pull Requests Releases Wiki Activity

updates

This commit is contained in:
Glenn Jocher 2020-06-04 14:25:10 -07:00
parent ba9ab663a2
commit a2336088f0
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
models/yolo.py
View File

@ -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()