From 09ee7b6f115d96a377558492e05f7ef703d1c390 Mon Sep 17 00:00:00 2001
From: Glenn Jocher <glenn.jocher@ultralytics.com>
Date: Fri, 3 May 2019 20:50:44 +0200
Subject: [PATCH] updates

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

diff --git a/models.py b/models.py
index dc008ada..33176bb9 100755
--- a/models.py
+++ b/models.py
@@ -165,6 +165,8 @@ class YOLOLayer(nn.Module):
             io[..., 4:] = torch.sigmoid(io[..., 4:])  # p_conf, p_cls
             # io[..., 5:] = F.softmax(io[..., 5:], dim=4)  # p_cls
             io[..., :4] *= self.stride
+            if self.nc == 1:  # single-class model https://github.com/ultralytics/yolov3/issues/235
+                io[..., 5] = 1
 
             # reshape from [1, 3, 13, 13, 85] to [1, 507, 85]
             return io.view(bs, -1, 5 + self.nc), p