Make LaProp weight decay match typical PyTorch 'decoupled' behaviour …

…where it's scaled by LR
huggingface · Nov 30, 2024 · 82e8677 · 82e8677
1 parent 886eb77
commit 82e8677
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/timm/optim/laprop.py b/timm/optim/laprop.py
@@ -116,6 +116,6 @@ def step(self, closure=None):
                 p.add_(exp_avg, alpha=-step_size)
 
                 if group['weight_decay'] != 0:
-                    p.add_(p, alpha=-group['weight_decay'])
+                    p.add_(p, alpha=-(group['lr'] * group['weight_decay']))
 
         return loss