This commit is contained in:
Glenn Jocher 2020-03-09 10:43:49 -07:00
parent 4a90221e79
commit 071d4113f6
1 changed files with 2 additions and 2 deletions

View File

@ -191,8 +191,8 @@ class YOLOLayer(nn.Module):
if ONNX_EXPORT: # grids must be computed in __init__ if ONNX_EXPORT: # grids must be computed in __init__
stride = [32, 16, 8][yolo_index] # stride of this layer stride = [32, 16, 8][yolo_index] # stride of this layer
nx = int(img_size[1] / stride) # number x grid points nx = img_size[1] // stride # number x grid points
ny = int(img_size[0] / stride) # number y grid points ny = img_size[0] // stride # number y grid points
create_grids(self, img_size, (nx, ny)) create_grids(self, img_size, (nx, ny))
def forward(self, p, img_size, var=None): def forward(self, p, img_size, var=None):