glenn-jocher commited on
Commit
23d2bde
·
1 Parent(s): bcd452c

leaf variable --single-cls training bug fix (#593)

Browse files
Files changed (1) hide show
  1. utils/utils.py +2 -1
utils/utils.py CHANGED
@@ -493,7 +493,8 @@ def compute_loss(p, targets, model): # predictions, targets, model
493
  s = 3 / np # output count scaling
494
  lbox *= h['giou'] * s
495
  lobj *= h['obj'] * s * (1.4 if np == 4 else 1.)
496
- lcls *= h['cls'] * s
 
497
  bs = tobj.shape[0] # batch size
498
 
499
  loss = lbox + lobj + lcls
 
493
  s = 3 / np # output count scaling
494
  lbox *= h['giou'] * s
495
  lobj *= h['obj'] * s * (1.4 if np == 4 else 1.)
496
+ if model.nc > 1:
497
+ lcls *= h['cls'] * s
498
  bs = tobj.shape[0] # batch size
499
 
500
  loss = lbox + lobj + lcls