foreach version actually works now..

lucidrains · Jul 31, 2024 · 32062de · 32062de
1 parent 5455604
commit 32062de
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 5 deletions.
diff --git a/adam_atan2_pytorch/adam_atan2.py b/adam_atan2_pytorch/adam_atan2.py
@@ -86,8 +86,8 @@ def step(
                 # using atan2 instead of a division with epsilon in denominator
                 # a * atan2(exp_avg / bias_correct1, b * sqrt(exp_avg_sq / bias_correct2))
 
-                den = exp_avg_sq.clone().mul_(b * b / bias_correct2).sqrt_()
-                update = exp_avg.clone().mul_(1. / bias_correct1).atan2_(den)
+                den = exp_avg_sq.mul(b * b / bias_correct2).sqrt_()
+                update = exp_avg.mul(1. / bias_correct1).atan2_(den)
 
                 # update parameters
 

diff --git a/adam_atan2_pytorch/foreach.py b/adam_atan2_pytorch/foreach.py
@@ -72,7 +72,7 @@ def step(
 
         for group in self.param_groups:
 
-            wd, lr, beta1, beta2, a, b = group['lr'], group['weight_decay'], *group['betas'], group['a'], group['b']
+            wd, lr, beta1, beta2, a, b = group['weight_decay'], group['lr'], *group['betas'], group['a'], group['b']
 
             # accumulate List[Tensor] for foreach inplace updates
 
@@ -123,7 +123,8 @@ def step(
 
             # weight decay
 
-            torch._foreach_mul_(params, 1. - lr * wd)
+            if wd > 0.:
+                torch._foreach_mul_(params, 1. - lr * wd)
 
             # decay running averages
 

diff --git a/pyproject.toml b/pyproject.toml
@@ -1,6 +1,6 @@
 [project]
 name = "adam-atan2-pytorch"
-version = "0.0.8"
+version = "0.0.9"
 description = "Adam-atan2 for Pytorch"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }