oops

Author: lucidrains

lion_pytorch/foreach.py

@@ -86,7 +86,7 @@ def step(
             torch._foreach_lerp_(updates, grads, 1. - beta1)
             torch._foreach_sign_(updates)
 
-            torch._foreach_add_(params, updates)
+            torch._foreach_add_(params, updates, alpha = -lr)
 
             # decay momentum running average
 

setup.py

@@ -3,7 +3,7 @@
 setup(
   name = 'lion-pytorch',
   packages = find_packages(exclude=[]),
-  version = '0.2.1',
+  version = '0.2.2',
   license='MIT',
   description = 'Lion Optimizer - Pytorch',
   author = 'Phil Wang',