diff --git a/models.py b/models.py
index 17ff234f..6159d776 100755
--- a/models.py
+++ b/models.py
@@ -185,18 +185,19 @@ class YOLOLayer(nn.Module):
 
         elif ONNX_EXPORT:
             # Constants CAN NOT BE BROADCAST, ensure correct shape!
-            ngu = self.ng.repeat((1, self.na * self.nx * self.ny, 1))
-            grid_xy = self.grid_xy.repeat((1, self.na, 1, 1, 1)).view((1, -1, 2))
-            anchor_wh = self.anchor_wh.repeat((1, 1, self.nx, self.ny, 1)).view((1, -1, 2)) / ngu
+            m = self.na * self.nx * self.ny
+            ngu = self.ng.repeat((1, m, 1))
+            grid_xy = self.grid_xy.repeat((1, self.na, 1, 1, 1)).view(1, m, 2)
+            anchor_wh = self.anchor_wh.repeat((1, 1, self.nx, self.ny, 1)).view(1, m, 2) / ngu
 
-            p = p.view(-1, 5 + self.nc)
+            p = p.view(m, 5 + self.nc)
             xy = torch.sigmoid(p[..., 0:2]) + grid_xy[0]  # x, y
             wh = torch.exp(p[..., 2:4]) * anchor_wh[0]  # width, height
             p_conf = torch.sigmoid(p[:, 4:5])  # Conf
             p_cls = F.softmax(p[:, 5:85], 1) * p_conf  # SSD-like conf
             return torch.cat((xy / ngu[0], wh, p_conf, p_cls), 1).t()
 
-            # p = p.view(1, -1, 5 + self.nc)
+            # p = p.view(1, m, 5 + self.nc)
             # xy = torch.sigmoid(p[..., 0:2]) + grid_xy  # x, y
             # wh = torch.exp(p[..., 2:4]) * anchor_wh  # width, height
             # p_conf = torch.sigmoid(p[..., 4:5])  # Conf
@@ -278,7 +279,7 @@ class Darknet(nn.Module):
         elif ONNX_EXPORT:
             output = torch.cat(output, 1)  # cat 3 layers 85 x (507, 2028, 8112) to 85 x 10647
             nc = self.module_list[self.yolo_layers[0]].nc  # number of classes
-            return output[5:5 + nc].t(), output[:4].t()  # ONNX scores, boxes
+            return output[5:5 + nc].t(), output[0:4].t()  # ONNX scores, boxes
         else:
             io, p = list(zip(*output))  # inference output, training output
             return torch.cat(io, 1), p