fix foreach

adam_atan2_pytorch/foreach.py

@@ -70,6 +70,14 @@ def step(
                 loss = closure()
 
         for group in self.param_groups:
+
+            # accumulate List[Tensor] for foreach inplace updates
+
+            params = []
+            grads = []
+            exp_avgs = []
+            exp_avg_sqs = []
+
             for p in filter(lambda p: exists(p.grad), group['params']):
 
                 grad, lr, wd, beta1, beta2, a, b, state, init_lr = p.grad, group['lr'], group['weight_decay'], *group['betas'], group['a'], group['b'], self.state[p], self._init_lr
@@ -79,13 +87,6 @@ def step(
                 if wd > 0.:
                     wd /= init_lr
 
-                # accumulate List[Tensor] for foreach inplace updates
-
-                params = []
-                grads = []
-                exp_avgs = []
-                exp_avg_sqs = []
-
                 # init state if needed
 
                 if len(state) == 0:

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "adam-atan2-pytorch"
-version = "0.0.5"
+version = "0.0.6"
 description = "Adam-atan2 for Pytorch"
 authors = [
     { name = "Phil Wang", email = "[email protected]" }