This commit is contained in:
Glenn Jocher 2019-12-08 18:30:10 -08:00
parent d603ac8e69
commit ca5da3dfe0
1 changed files with 1 additions and 3 deletions

View File

@ -167,9 +167,7 @@ class YOLOLayer(nn.Module):
create_grids(self, img_size, (nx, ny), p.device, p.dtype)
# p.view(bs, 255, 13, 13) -- > (bs, 3, 13, 13, 85) # (bs, anchors, grid, grid, classes + xywh)
# p = p.view(bs, self.na, self.nc + 5, self.ny, self.nx).permute(0, 1, 3, 4, 2).contiguous() # prediction
# https://discuss.pytorch.org/t/in-pytorch-0-4-is-it-recommended-to-use-reshape-than-view-when-it-is-possible/17034
p = p.reshape(bs, self.na, self.nc + 5, self.ny, self.nx).permute(0, 1, 3, 4, 2) # prediction
p = p.view(bs, self.na, self.nc + 5, self.ny, self.nx).permute(0, 1, 3, 4, 2).contiguous() # prediction
if self.training:
return p