This commit is contained in:
Glenn Jocher
2020-01-12 12:01:58 -08:00
parent aeac9b78eb
commit b890ccecfc
+1 -1
View File
@@ -206,7 +206,7 @@ class YOLOLayer(nn.Module):
if self.nc == 1: if self.nc == 1:
io[..., 5] = 1 # single-class model https://github.com/ultralytics/yolov3/issues/235 io[..., 5] = 1 # single-class model https://github.com/ultralytics/yolov3/issues/235
# reshape from [1, 3, 13, 13, 85] to [1, 507, 84], remove obj_conf # reshape from [1, 3, 13, 13, 85] to [1, 507, 85]
return io.view(bs, -1, self.no), p return io.view(bs, -1, self.no), p