From 76be876739c15f527aff67c2f0f379483c698e5f Mon Sep 17 00:00:00 2001
From: Ferdinand Schlatt <ferdinand.schlatt@uni-jena.de>
Date: Thu, 21 Nov 2024 13:59:25 +0100
Subject: [PATCH] allow other lr schedulers

---
 lightning_ir/lightning_utils/lr_schedulers.py | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/lightning_ir/lightning_utils/lr_schedulers.py b/lightning_ir/lightning_utils/lr_schedulers.py
index eed817c..6fec174 100644
--- 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,
+)