Merge pull request #227 from i404788/main

[Bugfix] Fix SRMM to allow operation beyond memory_length
kozistr · Apr 7, 2024 · 523f140 · 523f140
2 parents da65344 + b0124e0
commit 523f140
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/pytorch_optimizer/optimizer/srmm.py b/pytorch_optimizer/optimizer/srmm.py
@@ -53,7 +53,7 @@ def step(self, closure: CLOSURE = None) -> LOSS:
                 group['step'] = 1
 
             w_t: float = (
-                (group['step'] + 1) % (group['memory_length'] if group['memory_length'] is not None else 1)
+                (group['step'] % (group['memory_length'] if group['memory_length'] is not None else 1)) + 1
             ) ** -group['beta']
 
             for p in group['params']: