Add AutoWU LR scheduler from KakaoBrain

Signed-off-by: Emanuele Ballarin <emanuele@ballarin.cc>
emaballarin · Jul 23, 2023 · 647ed0d · 647ed0d
1 parent 441581b
commit 647ed0d
Show file tree

Hide file tree

Showing 5 changed files with 527 additions and 1 deletion.
diff --git a/ebtorch/__init__.py b/ebtorch/__init__.py
@@ -69,6 +69,7 @@
 del AdaHessian
 del Adan
 del Lion
+del AutoWU
 del Lookahead
 del Lookaround
 del RAdam

diff --git a/ebtorch/optim/__init__.py b/ebtorch/optim/__init__.py
@@ -24,6 +24,7 @@
 from .adabound import AdaBound
 from .adahessian import AdaHessian
 from .adan import Adan
+from .autowu import AutoWU
 from .custom import ralah_optim
 from .custom import wfneal
 from .lion import Lion
@@ -37,6 +38,7 @@
 del adan
 del custom
 del lion
+del autowu
 del lookahead
 del lookaround
 del radam