bug fix - not calculate cls_loss when class 1

Author: 2damin

train/loss.py

@@ -50,7 +50,7 @@ def compute_loss(self, pred, targets = None, yolo_layers = None, tmp_img = None)
 
                 tobj[batch_id, anc_id, gy, gx] = iou.detach().clamp(0).type(tobj.dtype)
 
-                if ps.size(1) - 5 > 1:
+                if ps.size(1) - 5 >= 1:
                     #all class loss
                     t = torch.zeros_like(ps[...,5:], device=self.device)
                     t[range(num_targets), tcls[pidx]] = 1