Fix weight decay comment (#4228)

ultralytics · Jul 29, 2021 · 8d3c3ef · 8d3c3ef
1 parent 1a10b0e
commit 8d3c3ef
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/train.py b/train.py
@@ -128,9 +128,9 @@ def train(hyp,  # path/to/hyp.yaml or hyp dictionary
     for v in model.modules():
         if hasattr(v, 'bias') and isinstance(v.bias, nn.Parameter):  # bias
             g2.append(v.bias)
-        if isinstance(v, nn.BatchNorm2d):  # weight with decay
+        if isinstance(v, nn.BatchNorm2d):  # weight (no decay)
             g0.append(v.weight)
-        elif hasattr(v, 'weight') and isinstance(v.weight, nn.Parameter):  # weight without decay
+        elif hasattr(v, 'weight') and isinstance(v.weight, nn.Parameter):  # weight (with decay)
             g1.append(v.weight)
 
     if opt.adam: