allow other lr schedulers

webis-de · Nov 21, 2024 · 76be876 · 76be876
1 parent e84281c
commit 76be876
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/lightning_ir/lightning_utils/lr_schedulers.py b/lightning_ir/lightning_utils/lr_schedulers.py
@@ -33,4 +33,9 @@ class ConstantLRSchedulerWithLinearWarmup(WarmupLRScheduler, ConstantSchedulerWi
     pass
 
 
-LR_SCHEDULERS = [LinearLRSchedulerWithLinearWarmup, ConstantLRSchedulerWithLinearWarmup]
+LR_SCHEDULERS = (
+    LinearLRSchedulerWithLinearWarmup,
+    ConstantLRSchedulerWithLinearWarmup,
+    WarmupLRScheduler,
+    torch.optim.lr_scheduler.LRScheduler,
+)