From d603ac8e694a3d7f1d8c777a47c5f049917e4b7e Mon Sep 17 00:00:00 2001
From: Glenn Jocher <glenn.jocher@ultralytics.com>
Date: Sun, 8 Dec 2019 18:08:19 -0800
Subject: [PATCH] updates

---
 models.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/models.py b/models.py
index 528ede98..cf95a959 100755
--- 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)