From e27b124828642198581512d42b14f0afe181ecd5 Mon Sep 17 00:00:00 2001
From: Glenn Jocher <glenn.jocher@ultralytics.com>
Date: Wed, 4 Dec 2019 17:50:52 -0800
Subject: [PATCH] updates

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

diff --git a/models.py b/models.py
index 6159d776..7cc82b58 100755
--- a/models.py
+++ b/models.py
@@ -194,7 +194,7 @@ class YOLOLayer(nn.Module):
             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
+            p_cls = F.softmax(p[:, 5:5 + self.nc], 1) * p_conf  # SSD-like conf
             return torch.cat((xy / ngu[0], wh, p_conf, p_cls), 1).t()
 
             # p = p.view(1, m, 5 + self.nc)