diff --git a/models.py b/models.py
index 0e92c79f..71453c23 100755
--- a/models.py
+++ b/models.py
@@ -137,17 +137,17 @@ class YOLOLayer(nn.Module):
             anchor_wh = self.anchor_wh.repeat((1, 1, nG, nG, 1)).view((1, -1, 2)) / nG
 
             # p = p.view(-1, 5 + self.nC)
-            # xy = xy + self.grid_xy[0]  # x, y
-            # wh = torch.exp(wh) * self.anchor_wh[0]  # width, height
+            # 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:], 1) * p_conf  # SSD-like conf
+            # p_cls = F.softmax(p[:, 5:85], 1) * p_conf  # SSD-like conf
             # return torch.cat((xy / nG, wh, p_conf, p_cls), 1).t()
 
             p = p.view(1, -1, 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
-            p_cls = p[..., 5:]
+            p_cls = p[..., 5:85]
             # Broadcasting only supported on first dimension in CoreML. See onnx-coreml/_operators.py
             # p_cls = F.softmax(p_cls, 2) * p_conf  # SSD-like conf
             p_cls = torch.exp(p_cls).permute((2, 1, 0))
@@ -203,8 +203,8 @@ class Darknet(nn.Module):
             layer_outputs.append(x)
 
         if ONNX_EXPORT:
-            output = torch.cat(output, 1)  # merge the 3 layers 85 x (507, 2028, 8112) to 85 x 10647
-            return output[5:].t(), output[:4].t()  # ONNX scores, boxes
+            output = torch.cat(output, 1)  # cat 3 layers 85 x (507, 2028, 8112) to 85 x 10647
+            return output[5:85].t(), output[:4].t()  # ONNX scores, boxes
         else:
             return output if self.training else torch.cat(output, 1)