diff --git a/models.py b/models.py index 204cc990..e6832ab5 100755 --- a/models.py +++ b/models.py @@ -159,7 +159,8 @@ class YOLOLayer(nn.Module): # Compute losses nT = sum([len(x) for x in targets]) # number of targets nM = mask.sum().float() # number of anchors (assigned to targets) - k = nM / bs + k = 1 #nM / bs + if nM > 0: lxy = k * MSELoss(xy[mask], txy[mask]) lwh = k * MSELoss(wh[mask], twh[mask])