updates

2019-12-08 18:08:19 -08:00 · 2019-12-08 18:08:19 -08:00 · d603ac8e69
commit d603ac8e69
parent 61c3cb9ecf
1 changed files with 1 additions and 1 deletions
--- a/models.py
+++ b/models.py
@ -162,7 +162,7 @@ class YOLOLayer(nn.Module):
        if ONNX_EXPORT:
            bs = 1  # batch size
        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):
                create_grids(self, img_size, (nx, ny), p.device, p.dtype)