update: REX lr scheduler

Author: kozistr

pytorch_optimizer/__init__.py

@@ -26,6 +26,7 @@
 from pytorch_optimizer.lr_scheduler.experimental.deberta_v3_lr_scheduler import deberta_v3_large_lr_scheduler
 from pytorch_optimizer.lr_scheduler.linear_warmup import CosineScheduler, LinearScheduler, PolyScheduler
 from pytorch_optimizer.lr_scheduler.proportion import ProportionScheduler
+from pytorch_optimizer.lr_scheduler.rex import REXScheduler
 from pytorch_optimizer.optimizer.a2grad import A2Grad
 from pytorch_optimizer.optimizer.adabelief import AdaBelief
 from pytorch_optimizer.optimizer.adabound import AdaBound
@@ -195,6 +196,7 @@
     PolyScheduler,
     LinearScheduler,
     ProportionScheduler,
+    REXScheduler,
 ]
 LR_SCHEDULERS: Dict[str, SCHEDULER] = {
     str(lr_scheduler.__name__).lower(): lr_scheduler for lr_scheduler in LR_SCHEDULER_LIST

pytorch_optimizer/base/scheduler.py

@@ -76,7 +76,6 @@ def step(self):
 
         self.step_t += 1
 
-        # apply the lr to optimizer if it's provided
         if self.optimizer is not None:
             for param_group in self.optimizer.param_groups:
                 param_group['lr'] = value

tests/test_load_modules.py

@@ -42,7 +42,7 @@ def test_get_supported_optimizers():
 
 
 def test_get_supported_lr_schedulers():
-    assert len(get_supported_lr_schedulers()) == 10
+    assert len(get_supported_lr_schedulers()) == 11
 
 
 def test_get_supported_loss_functions():