updates
This commit is contained in:
parent
61c3cb9ecf
commit
d603ac8e69
|
@ -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)
|
||||||
|
|
||||||
|
|
Loading…
Reference in New Issue