diff --git a/CHANGELOG.md b/CHANGELOG.md
index f503ba2d42..599a658557 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -4,6 +4,13 @@ All notable changes to this project will be documented in this file.
 
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/).
 
+## [v.0.7.0.dev150]
+
+### Updated
+
+- Update padim code to work properly with lightning 2.4.0
+- Move unused legacy dependencies to extras
+
 ## [v0.7.0.dev143]
 
 ### Fixed
diff --git a/pyproject.toml b/pyproject.toml
index f36ea4996f..6a75f30502 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "anomalib-orobix"
-version = "0.7.0.dev143"
+version = "0.7.0.dev150"
 description = "Orobix anomalib fork"
 authors = [
     "Intel OpenVINO <help@openvino.intel.com>",
@@ -19,21 +19,23 @@ requires = ["poetry-core>=1.0.0"]
 build-backend = "poetry.core.masonry.api"
 
 [tool.poetry.dependencies]
-python = ">=3.9,<3.11"
+python = ">=3.10,<3.11"
 einops = "~0.6"
 kornia = "0.6.5"
 omegaconf = "~2.3"
 freia = "~0.2"
-wandb = "0.12.17"
-gradio = "3.0.2"
+
 line-profiler = "3.5.1"
 jsonargparse = { version = "~4.3.0", extras = ["signatures"] }
 
 imgaug = { version = "0.4.0", optional = true }
-
+gradio = { version = "3.0.2", optional = true }
+wandb = { version = "0.12.17", optional = true }
 [tool.poetry.extras]
 # We make imgaug optional as it requires opencv-python instead of the headless version
 augmentation = ["imgaug"]
+ui = ["gradio"]
+wandb = ["wandb"]
 
 # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
 # BLACK CONFIGURATION                                                         #
diff --git a/src/anomalib/__init__.py b/src/anomalib/__init__.py
index 600adfd2bb..fcb74a8812 100644
--- a/src/anomalib/__init__.py
+++ b/src/anomalib/__init__.py
@@ -4,6 +4,6 @@
 # SPDX-License-Identifier: Apache-2.0
 
 anomalib_version = "0.7.0"
-custom_orobix_version = "1.4.3"
+custom_orobix_version = "1.5.0"
 
 __version__ = f"{anomalib_version}.dev{custom_orobix_version.replace('.', '')}"
diff --git a/src/anomalib/models/padim/lightning_model.py b/src/anomalib/models/padim/lightning_model.py
index 67c9f590a3..2047744b33 100644
--- a/src/anomalib/models/padim/lightning_model.py
+++ b/src/anomalib/models/padim/lightning_model.py
@@ -72,6 +72,7 @@ def configure_optimizers() -> None:  # pylint: disable=arguments-differ
     def on_train_epoch_start(self) -> None:
         self.embeddings = []
         self.stats = []
+        self.model.in_training_loop = True
         return super().on_train_epoch_start()
 
     def training_step(self, batch: dict[str, str | Tensor], *args, **kwargs) -> None:
@@ -109,6 +110,7 @@ def on_validation_start(self) -> None:
             logger.warning("No embeddings were extracted from the training set. Skipping Gaussian fitting.")
             return
 
+        self.model.in_training_loop = False
         logger.info("Aggregating the embedding extracted from the training set.")
 
         embeddings = torch.vstack(self.embeddings)
@@ -134,7 +136,14 @@ def validation_step(self, batch: dict[str, str | Tensor], *args, **kwargs) -> ST
         batch["anomaly_maps"], _ = self.model(batch["image"])
         return batch
 
+    def on_test_start(self):
+        self.model.in_training_loop = False
+        return super().on_test_start()
 
+    def on_predict_start(self):
+        self.model.in_training_loop = False
+        return super().on_predict_start()
+    
 class PadimLightning(Padim):
     """PaDiM: a Patch Distribution Modeling Framework for Anomaly Detection and Localization.
 
diff --git a/src/anomalib/models/padim/torch_model.py b/src/anomalib/models/padim/torch_model.py
index 6b611e0b90..b0c1bc63af 100644
--- a/src/anomalib/models/padim/torch_model.py
+++ b/src/anomalib/models/padim/torch_model.py
@@ -121,6 +121,7 @@ def __init__(
         self.anomaly_map_generator = AnomalyMapGenerator(image_size=input_size)
 
         self.gaussian = MultiVariateGaussian(self.n_features, self.n_patches, tied_covariance=tied_covariance)
+        self.in_training_loop = False
 
     def forward(self, input_tensor: Tensor) -> Tuple[Tensor, Tensor]:
         """Forward-pass image-batch (N, C, H, W) into model to extract features.
@@ -160,7 +161,7 @@ def forward(self, input_tensor: Tensor) -> Tuple[Tensor, Tensor]:
 
         anomaly_score = None
 
-        if self.training:
+        if self.in_training_loop:
             output = embeddings
         else:
             output = self.anomaly_map_generator(