works fine for me

lucidrains · lucidrains · commit 110ea77ff9f5 · 2024-10-23T07:37:46.000-07:00
diff --git a/adam_atan2_pytorch/__init__.py b/adam_atan2_pytorch/__init__.py
@@ -1 +1,3 @@
 from adam_atan2_pytorch.adam_atan2 import AdamAtan2
+
+Adam = AdamAtan2
diff --git a/adam_atan2_pytorch/adam_atan2_with_wasserstein_reg.py b/adam_atan2_pytorch/adam_atan2_with_wasserstein_reg.py
@@ -137,3 +137,5 @@ def step(
                 state['steps'] = steps
 
         return loss
+
+Adam = AdamAtan2
diff --git a/adam_atan2_pytorch/foreach.py b/adam_atan2_pytorch/foreach.py
@@ -168,3 +168,5 @@ def step(
             torch._foreach_add_(params, updates, alpha = -lr * a)
 
         return loss
+
+Adam = AdamAtan2

Original file line number	Diff line number	Diff line change
`@@ -1 +1,3 @@`
`1`	`1`	`from adam_atan2_pytorch.adam_atan2 import AdamAtan2`
	`2`	`+`
	`3`	`+Adam = AdamAtan2`