Bump lightning version to 2.3.3 (#3796)

* update for resolving cve-2024-5980 from lightning 2.2.2 * update torch lr scheduler importing logic
openvinotoolkit · Aug 8, 2024 · c61e7d6 · c61e7d6
1 parent 63842a9
commit c61e7d6
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 5 deletions.
diff --git a/pyproject.toml b/pyproject.toml
@@ -92,7 +92,7 @@ xpu = [
 
 base = [
     "torch==2.2.2",
-    "lightning==2.2",
+    "lightning==2.3.3",
     "pytorchcv",
     "timm==1.0.3",
     "openvino==2024.3",

diff --git a/src/otx/core/schedulers/__init__.py b/src/otx/core/schedulers/__init__.py
@@ -7,19 +7,18 @@
 
 from typing import Callable
 
-from lightning.fabric.utilities.types import _TORCH_LRSCHEDULER
 from lightning.pytorch.cli import ReduceLROnPlateau
+from torch.optim.lr_scheduler import LRScheduler
 from torch.optim.optimizer import Optimizer
 
 from otx.core.schedulers.callable import SchedulerCallableSupportHPO
 from otx.core.schedulers.warmup_schedulers import LinearWarmupScheduler, LinearWarmupSchedulerCallable
 
+LRSchedulerListCallable = Callable[[Optimizer], list[LRScheduler | ReduceLROnPlateau]]
+
 __all__ = [
     "LRSchedulerListCallable",
     "LinearWarmupScheduler",
     "LinearWarmupSchedulerCallable",
     "SchedulerCallableSupportHPO",
 ]
-
-
-LRSchedulerListCallable = Callable[[Optimizer], list[_TORCH_LRSCHEDULER | ReduceLROnPlateau]]