Fiddling with Kron (PSGD)

Author: rwightman

timm/optim/_optim_factory.py

@@ -23,6 +23,7 @@
 from .adamw import AdamWLegacy
 from .adan import Adan
 from .adopt import Adopt
+from .kron import Kron
 from .lamb import Lamb
 from .laprop import LaProp
 from .lars import Lars
@@ -693,6 +694,12 @@ def _register_other_optimizers(registry: OptimizerRegistry) -> None:
             has_betas=True,
             second_order=True,
         ),
+        OptimInfo(
+            name='kron',
+            opt_class=Kron,
+            description='',
+            has_momentum=True,
+        ),
         OptimInfo(
             name='laprop',
             opt_class=LaProp,