cleanup, save experiments for sunday

lucidrains · lucidrains · commit 9f0fa124c5ee · 2024-06-15T12:49:53.000-07:00
diff --git a/grokfast_pytorch/grokfast.py b/grokfast_pytorch/grokfast.py
@@ -113,7 +113,7 @@ def step(
 
                     # appendix C - line 24 of https://arxiv.org/html/2405.20233v2
 
-                    grad.add_(grok_exp_avg * lamb)
+                    grad.add_(grok_exp_avg, alpha = lamb)
 
                 # bias corrections
 
