cleanup, readme, and release

Author: lucidrains

README.md

@@ -6,6 +6,8 @@ Implementation of the proposed <a href="https://arxiv.org/abs/2407.05872">Adam-a
 
 A multi-million dollar paper out of google deepmind proposes a small change to Adam update rule (using `atan2`) to remove the epsilon altogether for numerical stability and scale invariance
 
+It also contains some features for improving plasticity (continual learning field)
+
 ## Install
 
 ```bash

adam_atan2_pytorch/adam_atan2.py

@@ -20,7 +20,6 @@ def __init__(
         betas: tuple[float, float] = (0.9, 0.99),
         weight_decay = 0.,
         regen_reg_rate = 0.,
-        wasserstein_reg = False,
         decoupled_wd = False,
         a = 1.27,
         b = 1.
@@ -41,7 +40,6 @@ def __init__(
             b = b,
             weight_decay = weight_decay,
             regen_reg_rate = regen_reg_rate,
-            wasserstein_reg = wasserstein_reg,
         )
 
         super().__init__(params, defaults)
@@ -60,7 +58,7 @@ def step(
         for group in self.param_groups:
             for p in filter(lambda p: exists(p.grad), group['params']):
 
-                grad, lr, wd, regen_rate, wasserstein_reg, beta1, beta2, a, b, state, init_lr = p.grad, group['lr'], group['weight_decay'], group['regen_reg_rate'], group['wasserstein_reg'], *group['betas'], group['a'], group['b'], self.state[p], self._init_lr
+                grad, lr, wd, regen_rate, beta1, beta2, a, b, state, init_lr = p.grad, group['lr'], group['weight_decay'], group['regen_reg_rate'], *group['betas'], group['a'], group['b'], self.state[p], self._init_lr
 
                 # maybe decoupled weight decay
 

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "adam-atan2-pytorch"
-version = "0.0.15"
+version = "0.1.0"
 description = "Adam-atan2 for Pytorch"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }