From 3ba91b1e53a7b55ecb8fb39f3722868d32661c20 Mon Sep 17 00:00:00 2001
From: Ferdinand Schlatt <ferdinand.schlatt@uni-jena.de>
Date: Thu, 21 Nov 2024 14:00:08 +0100
Subject: [PATCH] remove lr_schedulers from main module

---
 lightning_ir/__init__.py                 | 2 --
 lightning_ir/lightning_utils/__init__.py | 8 +-------
 2 files changed, 1 insertion(+), 9 deletions(-)

diff --git a/lightning_ir/__init__.py b/lightning_ir/__init__.py
index 14fe378..780d2b8 100644
--- a/lightning_ir/__init__.py
+++ b/lightning_ir/__init__.py
@@ -42,7 +42,6 @@
     TupleDataset,
 )
 from .lightning_utils import (
-    LR_SCHEDULERS,
     ConstantLRSchedulerWithLinearWarmup,
     GenericConstantSchedulerWithLinearWarmup,
     GenericConstantSchedulerWithQuadraticWarmup,
@@ -180,7 +179,6 @@
     "LightningIRWandbLogger",
     "LinearLRSchedulerWithLinearWarmup",
     "LocalizedContrastiveEstimation",
-    "LR_SCHEDULERS",
     "QueryDataset",
     "QuerySample",
     "RankCallback",
diff --git a/lightning_ir/lightning_utils/__init__.py b/lightning_ir/lightning_utils/__init__.py
index b2b6268..b7ccca7 100644
--- a/lightning_ir/lightning_utils/__init__.py
+++ b/lightning_ir/lightning_utils/__init__.py
@@ -1,10 +1,5 @@
 from .callbacks import IndexCallback, RankCallback, ReRankCallback, SearchCallback
-from .lr_schedulers import (
-    LR_SCHEDULERS,
-    ConstantLRSchedulerWithLinearWarmup,
-    LinearLRSchedulerWithLinearWarmup,
-    WarmupLRScheduler,
-)
+from .lr_schedulers import ConstantLRSchedulerWithLinearWarmup, LinearLRSchedulerWithLinearWarmup, WarmupLRScheduler
 from .schedulers import (
     GenericConstantSchedulerWithLinearWarmup,
     GenericConstantSchedulerWithQuadraticWarmup,
@@ -18,7 +13,6 @@
     "GenericLinearSchedulerWithLinearWarmup",
     "IndexCallback",
     "LinearLRSchedulerWithLinearWarmup",
-    "LR_SCHEDULERS",
     "RankCallback",
     "ReRankCallback",
     "SearchCallback",