From c61e7d6d60a28403e73474aa64ce8c8452341b02 Mon Sep 17 00:00:00 2001
From: Yunchu Lee <yunchu.lee@intel.com>
Date: Thu, 8 Aug 2024 09:40:59 +0900
Subject: [PATCH] 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
---
 pyproject.toml                      | 2 +-
 src/otx/core/schedulers/__init__.py | 7 +++----
 2 files changed, 4 insertions(+), 5 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index e913a9df994..c6aaef76da3 100644
--- 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
index ed74aaf62cf..d4d928e3c61 100644
--- 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]]