use inplace sign

lucidrains · Feb 17, 2023 · 5d63c0e · 5d63c0e
1 parent b67f8d8
commit 5d63c0e
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/lion_pytorch/lion_pytorch.py b/lion_pytorch/lion_pytorch.py
@@ -17,8 +17,8 @@ def update_fn(p, grad, exp_avg, lr, wd, beta1, beta2):
 
     # weight update
 
-    update = exp_avg.clone().lerp_(grad, 1 - beta1)
-    p.add_(torch.sign(update), alpha = -lr)
+    update = exp_avg.clone().lerp_(grad, 1 - beta1).sign_()
+    p.add_(update, alpha = -lr)
 
     # decay the momentum running average coefficient
 

diff --git a/setup.py b/setup.py
@@ -3,7 +3,7 @@
 setup(
   name = 'lion-pytorch',
   packages = find_packages(exclude=[]),
-  version = '0.0.4',
+  version = '0.0.5',
   license='MIT',
   description = 'Lion Optimizer - Pytorch',
   author = 'Phil Wang',