fix(optim): self-graft correctly

Author: ClashLuke

setup.py

@@ -10,7 +10,7 @@
     name='truegrad',
     license='BSD',
     description='PyTorch interface for TrueGrad-AdamW',
-    version='4.0.2',
+    version='4.0.3',
     long_description=README,
     url='https://github.com/clashluke/truegrad',
     packages=setuptools.find_packages(),
@@ -26,6 +26,7 @@
         'Programming Language :: Python :: 3.7',
         'Programming Language :: Python :: 3.8',
         'Programming Language :: Python :: 3.9',
+        'Programming Language :: Python :: 3.10',
         'Topic :: Software Development :: Libraries',
         'Topic :: Software Development :: Libraries :: Python Modules',
         'Intended Audience :: Developers',

truegrad/optim.py

@@ -196,9 +196,10 @@ def step(self, closure=None):
                 update = p.double() - o.double()
                 p.set_(o)
                 scale = group["lr"]
+                sign_update = torch.sign(update)
                 if group["graft_to_self"]:
-                    scale = scale * torch.norm(update)
-                p.add_(torch.sign(update), alpha=scale)
+                    scale = scale * update.norm() / sign_update.norm().clamp(min=group["eps"])
+                p.add_(sign_update, alpha=scale)
 
         return loss