Skip to content

Commit

Permalink
fix foreach
Browse files Browse the repository at this point in the history
  • Loading branch information
lucidrains committed Jul 31, 2024
1 parent ff2d25d commit 0044795
Show file tree
Hide file tree
Showing 2 changed files with 9 additions and 8 deletions.
15 changes: 8 additions & 7 deletions adam_atan2_pytorch/foreach.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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:
Expand Down
2 changes: 1 addition & 1 deletion pyproject.toml
Original file line number Diff line number Diff line change
@@ -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 = "lucidrains@gmail.com" }
Expand Down

0 comments on commit 0044795

Please sign in to comment.