apoorva/greenhouse
2
1
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

updates

Browse Source
This commit is contained in:
Glenn Jocher 2019-12-08 18:08:19 -08:00
parent 61c3cb9ecf
commit d603ac8e69
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.py
View File

@ -162,7 +162,7 @@ class YOLOLayer(nn.Module):
if ONNX_EXPORT: if ONNX_EXPORT:
bs = 1 # batch size bs = 1 # batch size
else: else:
bs, ny, nx = p.shape[0], p.shape[-2], p.shape[-1] bs, _, ny, nx = p.shape # bs, 255, 13, 13
if (self.nx, self.ny) != (nx, ny): if (self.nx, self.ny) != (nx, ny):
create_grids(self, img_size, (nx, ny), p.device, p.dtype) create_grids(self, img_size, (nx, ny), p.device, p.dtype)