From 925b91c20f86e9d715a825c1a81fa9492a11bb45 Mon Sep 17 00:00:00 2001
From: Weilin Xu <weilin.xu@intel.com>
Date: Tue, 27 May 2025 16:41:59 -0700
Subject: [PATCH 01/18] Detect if lightning is installed.

---
 mart/utils/imports.py    | 1 +
 tests/test_dependency.py | 7 ++++++-
 2 files changed, 7 insertions(+), 1 deletion(-)

diff --git a/mart/utils/imports.py b/mart/utils/imports.py
index d053ce5e..332170ea 100644
--- a/mart/utils/imports.py
+++ b/mart/utils/imports.py
@@ -25,3 +25,4 @@ def has(module_name):
 _HAS_TORCHVISION = has("torchvision")
 _HAS_TIMM = has("timm")
 _HAS_PYCOCOTOOLS = has("pycocotools")
+_HAS_LIGHTNING = has("lightning")
diff --git a/tests/test_dependency.py b/tests/test_dependency.py
index 5ee5ed24..4b7961b9 100644
--- a/tests/test_dependency.py
+++ b/tests/test_dependency.py
@@ -8,6 +8,7 @@
 
 from mart.utils.imports import (
     _HAS_FIFTYONE,
+    _HAS_LIGHTNING,
     _HAS_PYCOCOTOOLS,
     _HAS_TIMM,
     _HAS_TORCHVISION,
@@ -17,5 +18,9 @@
 def test_dependency_on_ci():
     if os.getenv("CI") == "true":
         assert (
-            _HAS_FIFTYONE and _HAS_TIMM and _HAS_PYCOCOTOOLS and _HAS_TORCHVISION is True
+            _HAS_FIFTYONE
+            and _HAS_TIMM
+            and _HAS_PYCOCOTOOLS
+            and _HAS_TORCHVISION
+            and _HAS_LIGHTNING is True
         ), "The dependency is not complete on CI, thus some tests are skipped."

From 14e8fdf6986b9d225c67d6eb8306dd375eca1c9d Mon Sep 17 00:00:00 2001
From: Weilin Xu <weilin.xu@intel.com>
Date: Tue, 27 May 2025 17:18:08 -0700
Subject: [PATCH 02/18] MisconfigurationException -> RuntimeError, because it
 is not misue of lightning.

---
 mart/attack/perturber.py | 9 +++------
 1 file changed, 3 insertions(+), 6 deletions(-)

diff --git a/mart/attack/perturber.py b/mart/attack/perturber.py
index 76ff1f42..493cf2e4 100644
--- a/mart/attack/perturber.py
+++ b/mart/attack/perturber.py
@@ -9,7 +9,6 @@
 from typing import TYPE_CHECKING, Callable, Iterable, Sequence
 
 import torch
-from lightning.pytorch.utilities.exceptions import MisconfigurationException
 
 from .projector import Projector
 
@@ -85,21 +84,19 @@ def create_from_tensor(tensor):
 
     def named_parameters(self, *args, **kwargs):
         if self.perturbation is None:
-            raise MisconfigurationException("You need to call configure_perturbation before fit.")
+            raise RuntimeError("You need to call configure_perturbation before fit.")
 
         return super().named_parameters(*args, **kwargs)
 
     def parameters(self, *args, **kwargs):
         if self.perturbation is None:
-            raise MisconfigurationException("You need to call configure_perturbation before fit.")
+            raise RuntimeError("You need to call configure_perturbation before fit.")
 
         return super().parameters(*args, **kwargs)
 
     def forward(self, **batch):
         if self.perturbation is None:
-            raise MisconfigurationException(
-                "You need to call the configure_perturbation before forward."
-            )
+            raise RuntimeError("You need to call the configure_perturbation before forward.")
 
         self.projector_(self.perturbation, **batch)
         # We need to register the hook at every forward pass.

From 027d4d0d3efe361ec83a87c2dbb48da1dd3b157b Mon Sep 17 00:00:00 2001
From: Weilin Xu <weilin.xu@intel.com>
Date: Tue, 27 May 2025 17:29:46 -0700
Subject: [PATCH 03/18] Conditional imports.

---
 mart/attack/__init__.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/mart/attack/__init__.py b/mart/attack/__init__.py
index 843ce9bd..6e62c9a3 100644
--- a/mart/attack/__init__.py
+++ b/mart/attack/__init__.py
@@ -1,4 +1,4 @@
-from .adversary import *
+from ..utils.imports import _HAS_LIGHTNING
 from .adversary_wrapper import *
 from .composer import *
 from .enforcer import *
@@ -8,3 +8,6 @@
 from .objective import *
 from .perturber import *
 from .projector import *
+
+if _HAS_LIGHTNING:
+    from .adversary import *

From b20acb54f663a38d2df7f97b08d646bfa0b79b0a Mon Sep 17 00:00:00 2001
From: Weilin Xu <weilin.xu@intel.com>
Date: Tue, 27 May 2025 17:32:01 -0700
Subject: [PATCH 04/18] Use the python logger instead in checking imports.

---
 mart/utils/imports.py | 10 ++++++----
 1 file changed, 6 insertions(+), 4 deletions(-)

diff --git a/mart/utils/imports.py b/mart/utils/imports.py
index 332170ea..01a82a8c 100644
--- a/mart/utils/imports.py
+++ b/mart/utils/imports.py
@@ -4,17 +4,19 @@
 # SPDX-License-Identifier: BSD-3-Clause
 #
 
+import logging
 from importlib.util import find_spec
 
-from .pylogger import get_pylogger
-
-logger = get_pylogger(__name__)
+# Avoid importing .pylogger when checking imports before running other code.
+logger = logging.getLogger(__name__)
 
 
 def has(module_name):
     module = find_spec(module_name)
     if module is None:
-        logger.warn(f"{module_name} is not installed, so some features in MART are unavailable.")
+        logger.warning(
+            f"{module_name} is not installed, so some features in MART are unavailable."
+        )
         return False
     else:
         return True

From e0bae506ace06547c71f6d45affb7850b9068084 Mon Sep 17 00:00:00 2001
From: Weilin Xu <weilin.xu@intel.com>
Date: Tue, 27 May 2025 17:33:19 -0700
Subject: [PATCH 05/18] Fallback to the basic python logger if lightning is not
 installed.

---
 mart/utils/pylogger.py | 13 +++++++++----
 1 file changed, 9 insertions(+), 4 deletions(-)

diff --git a/mart/utils/pylogger.py b/mart/utils/pylogger.py
index 86a723b9..d31023cf 100644
--- a/mart/utils/pylogger.py
+++ b/mart/utils/pylogger.py
@@ -1,6 +1,9 @@
 import logging
 
-from lightning.pytorch.utilities import rank_zero_only
+from .imports import _HAS_LIGHTNING
+
+if _HAS_LIGHTNING:
+    from lightning.pytorch.utilities import rank_zero_only
 
 __all__ = ["get_pylogger"]
 
@@ -12,8 +15,10 @@ def get_pylogger(name=__name__) -> logging.Logger:
 
     # this ensures all logging levels get marked with the rank zero decorator
     # otherwise logs would get multiplied for each GPU process in multi-GPU setup
-    logging_levels = ("debug", "info", "warning", "error", "exception", "fatal", "critical")
-    for level in logging_levels:
-        setattr(logger, level, rank_zero_only(getattr(logger, level)))
+    if _HAS_LIGHTNING:
+        logging_levels = ("debug", "info", "warning", "error", "exception", "fatal", "critical")
+        for level in logging_levels:
+            setattr(logger, level, rank_zero_only(getattr(logger, level)))
+    # Otherwise, fallback to the regular logger if lightning is not installed.
 
     return logger

From 9af8ba2b47617892b92bfeb392f866cac727ca6f Mon Sep 17 00:00:00 2001
From: Weilin Xu <weilin.xu@intel.com>
Date: Tue, 27 May 2025 17:41:06 -0700
Subject: [PATCH 06/18] Conditional imports in mart.utils.

---
 mart/utils/__init__.py | 11 +++++++----
 1 file changed, 7 insertions(+), 4 deletions(-)

diff --git a/mart/utils/__init__.py b/mart/utils/__init__.py
index 30e41c0a..2b7fff92 100644
--- a/mart/utils/__init__.py
+++ b/mart/utils/__init__.py
@@ -1,11 +1,14 @@
 from .adapters import *
-from .config import *
-from .imports import _HAS_TORCHVISION
+from .imports import _HAS_LIGHTNING, _HAS_TORCHVISION
 from .monkey_patch import *
 from .pylogger import *
-from .rich_utils import *
 from .silent import *
-from .utils import *
+
+if _HAS_LIGHTNING:
+    from .config import *
+    from .rich_utils import *
+    from .utils import *
+
 
 if _HAS_TORCHVISION:
     from .export import *

From 53da0168e3cf2ff405e3e7484014579c30ef86c2 Mon Sep 17 00:00:00 2001
From: Weilin Xu <weilin.xu@intel.com>
Date: Tue, 27 May 2025 17:43:14 -0700
Subject: [PATCH 07/18] Conditional imports in mart.__init__.

---
 mart/__init__.py | 9 ++++++---
 1 file changed, 6 insertions(+), 3 deletions(-)

diff --git a/mart/__init__.py b/mart/__init__.py
index 85181105..82d5d16b 100644
--- a/mart/__init__.py
+++ b/mart/__init__.py
@@ -1,11 +1,14 @@
-import importlib
+import importlib.metadata
 
 from mart import attack as attack
-from mart import datamodules as datamodules
-from mart import models as models
 from mart import nn as nn
 from mart import optim as optim
 from mart import transforms as transforms
 from mart import utils as utils
+from mart.utils.imports import _HAS_LIGHTNING
+
+if _HAS_LIGHTNING:
+    from mart import datamodules as datamodules
+    from mart import models as models
 
 __version__ = importlib.metadata.version(__package__ or __name__)

From dc94fef4862cec4701f5f93c120e2385cf664e0e Mon Sep 17 00:00:00 2001
From: Weilin Xu <weilin.xu@intel.com>
Date: Wed, 28 May 2025 21:48:01 -0700
Subject: [PATCH 08/18] Avoid importing all utils.

---
 mart/utils/__init__.py | 14 --------------
 1 file changed, 14 deletions(-)
 delete mode 100644 mart/utils/__init__.py

diff --git a/mart/utils/__init__.py b/mart/utils/__init__.py
deleted file mode 100644
index 2b7fff92..00000000
--- a/mart/utils/__init__.py
+++ /dev/null
@@ -1,14 +0,0 @@
-from .adapters import *
-from .imports import _HAS_LIGHTNING, _HAS_TORCHVISION
-from .monkey_patch import *
-from .pylogger import *
-from .silent import *
-
-if _HAS_LIGHTNING:
-    from .config import *
-    from .rich_utils import *
-    from .utils import *
-
-
-if _HAS_TORCHVISION:
-    from .export import *

From f523cb1b5aad2da7ab5e0403847bcd5901a1ac7a Mon Sep 17 00:00:00 2001
From: Weilin Xu <weilin.xu@intel.com>
Date: Mon, 2 Jun 2025 15:54:22 -0700
Subject: [PATCH 09/18] Separate lightning utils.

---
 mart/models/modular.py                | 2 +-
 mart/utils/{utils.py => lightning.py} | 0
 tests/test_utils.py                   | 2 +-
 3 files changed, 2 insertions(+), 2 deletions(-)
 rename mart/utils/{utils.py => lightning.py} (100%)

diff --git a/mart/models/modular.py b/mart/models/modular.py
index 67cb9d5e..ca10ab49 100644
--- a/mart/models/modular.py
+++ b/mart/models/modular.py
@@ -13,7 +13,7 @@
 
 from ..nn import SequentialDict
 from ..optim import OptimizerFactory
-from ..utils import flatten_dict
+from ..utils.lightning import flatten_dict
 from ..utils.optimization import configure_optimizers
 
 logger = logging.getLogger(__name__)
diff --git a/mart/utils/utils.py b/mart/utils/lightning.py
similarity index 100%
rename from mart/utils/utils.py
rename to mart/utils/lightning.py
diff --git a/tests/test_utils.py b/tests/test_utils.py
index cc1f6911..b2aa3b44 100644
--- a/tests/test_utils.py
+++ b/tests/test_utils.py
@@ -6,7 +6,7 @@
 
 import pytest
 
-from mart.utils import flatten_dict
+from mart.utils.lightning import flatten_dict
 
 
 def test_flatten_dict():

From 619af58b29e61e30fcba72f0656c58e7006db29e Mon Sep 17 00:00:00 2001
From: Weilin Xu <weilin.xu@intel.com>
Date: Mon, 2 Jun 2025 15:55:53 -0700
Subject: [PATCH 10/18] Only import utils that do not have external dependency.

---
 mart/callbacks/adversary_connector.py | 3 ++-
 mart/utils/__init__.py                | 5 +++++
 2 files changed, 7 insertions(+), 1 deletion(-)
 create mode 100644 mart/utils/__init__.py

diff --git a/mart/callbacks/adversary_connector.py b/mart/callbacks/adversary_connector.py
index 006fbe1f..d243da46 100644
--- a/mart/callbacks/adversary_connector.py
+++ b/mart/callbacks/adversary_connector.py
@@ -12,7 +12,8 @@
 import torch
 from lightning.pytorch.callbacks import Callback
 
-from ..utils import MonkeyPatch, get_pylogger
+from ..utils import MonkeyPatch
+from ..utils.pylogger import get_pylogger
 
 logger = get_pylogger(__name__)
 
diff --git a/mart/utils/__init__.py b/mart/utils/__init__.py
new file mode 100644
index 00000000..986c6374
--- /dev/null
+++ b/mart/utils/__init__.py
@@ -0,0 +1,5 @@
+# Only import components without external dependency.
+from .adapters import *
+from .monkey_patch import *
+from .optimization import *
+from .silent import *

From 5daf86773dd09a56cbe4a831a45fe59ba00d5861 Mon Sep 17 00:00:00 2001
From: Weilin Xu <weilin.xu@intel.com>
Date: Mon, 2 Jun 2025 15:56:29 -0700
Subject: [PATCH 11/18] Fix test.

---
 tests/test_perturber.py | 9 ++-------
 1 file changed, 2 insertions(+), 7 deletions(-)

diff --git a/tests/test_perturber.py b/tests/test_perturber.py
index af627359..cc7cdf64 100644
--- a/tests/test_perturber.py
+++ b/tests/test_perturber.py
@@ -4,15 +4,10 @@
 # SPDX-License-Identifier: BSD-3-Clause
 #
 
-from functools import partial
 from unittest.mock import Mock
 
 import pytest
-import torch
-from lightning.pytorch.utilities.exceptions import MisconfigurationException
-from torch.optim import SGD
 
-import mart
 from mart.attack import Perturber
 
 
@@ -36,10 +31,10 @@ def test_misconfiguration(input_data, target_data):
 
     perturber = Perturber(initializer=initializer, projector=projector)
 
-    with pytest.raises(MisconfigurationException):
+    with pytest.raises(RuntimeError):
         perturber(input=input_data, target=target_data)
 
-    with pytest.raises(MisconfigurationException):
+    with pytest.raises(RuntimeError):
         perturber.parameters()
 
 

From f9ec1ce020e5879437b4b1dffac3a293eabca61f Mon Sep 17 00:00:00 2001
From: Weilin Xu <weilin.xu@intel.com>
Date: Mon, 2 Jun 2025 16:11:04 -0700
Subject: [PATCH 12/18] Fix imports from mart.utils.lightning.

---
 mart/__main__.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/mart/__main__.py b/mart/__main__.py
index 9c44ee60..7d6167bf 100644
--- a/mart/__main__.py
+++ b/mart/__main__.py
@@ -11,7 +11,7 @@
 
 from mart import utils
 
-log = utils.get_pylogger(__name__)
+log = utils.pylogger.get_pylogger(__name__)
 
 # project root setup
 # uses the current working directory as root.
@@ -46,7 +46,7 @@ def main(cfg: DictConfig) -> float:
     # imports can be nested inside @hydra.main to optimize tab completion
     # https://github.com/facebookresearch/hydra/issues/934
     from mart.tasks.lightning import lightning
-    from mart.utils import get_metric_value, get_resume_checkpoint
+    from mart.utils.lightning import get_metric_value, get_resume_checkpoint
 
     # Resume and modify configs at the earliest point.
     # The actual checkpoint path is in cfg.ckpt_path

From f27a9d53d2659465dc6314819fb7f6abeb031fc7 Mon Sep 17 00:00:00 2001
From: Weilin Xu <weilin.xu@intel.com>
Date: Mon, 2 Jun 2025 16:37:13 -0700
Subject: [PATCH 13/18] Revert "Fallback to the basic python logger if
 lightning is not installed."

This reverts commit e0bae506ace06547c71f6d45affb7850b9068084.
---
 mart/utils/pylogger.py | 13 ++++---------
 1 file changed, 4 insertions(+), 9 deletions(-)

diff --git a/mart/utils/pylogger.py b/mart/utils/pylogger.py
index d31023cf..86a723b9 100644
--- a/mart/utils/pylogger.py
+++ b/mart/utils/pylogger.py
@@ -1,9 +1,6 @@
 import logging
 
-from .imports import _HAS_LIGHTNING
-
-if _HAS_LIGHTNING:
-    from lightning.pytorch.utilities import rank_zero_only
+from lightning.pytorch.utilities import rank_zero_only
 
 __all__ = ["get_pylogger"]
 
@@ -15,10 +12,8 @@ def get_pylogger(name=__name__) -> logging.Logger:
 
     # this ensures all logging levels get marked with the rank zero decorator
     # otherwise logs would get multiplied for each GPU process in multi-GPU setup
-    if _HAS_LIGHTNING:
-        logging_levels = ("debug", "info", "warning", "error", "exception", "fatal", "critical")
-        for level in logging_levels:
-            setattr(logger, level, rank_zero_only(getattr(logger, level)))
-    # Otherwise, fallback to the regular logger if lightning is not installed.
+    logging_levels = ("debug", "info", "warning", "error", "exception", "fatal", "critical")
+    for level in logging_levels:
+        setattr(logger, level, rank_zero_only(getattr(logger, level)))
 
     return logger

From 2b84ac6ecc464f68b8c59d4be591a237b797b0d7 Mon Sep 17 00:00:00 2001
From: Weilin Xu <weilin.xu@intel.com>
Date: Mon, 2 Jun 2025 16:40:27 -0700
Subject: [PATCH 14/18] Revert "Fix imports from mart.utils.lightning."

This reverts commit f9ec1ce020e5879437b4b1dffac3a293eabca61f.
---
 mart/__main__.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/mart/__main__.py b/mart/__main__.py
index 7d6167bf..9c44ee60 100644
--- a/mart/__main__.py
+++ b/mart/__main__.py
@@ -11,7 +11,7 @@
 
 from mart import utils
 
-log = utils.pylogger.get_pylogger(__name__)
+log = utils.get_pylogger(__name__)
 
 # project root setup
 # uses the current working directory as root.
@@ -46,7 +46,7 @@ def main(cfg: DictConfig) -> float:
     # imports can be nested inside @hydra.main to optimize tab completion
     # https://github.com/facebookresearch/hydra/issues/934
     from mart.tasks.lightning import lightning
-    from mart.utils.lightning import get_metric_value, get_resume_checkpoint
+    from mart.utils import get_metric_value, get_resume_checkpoint
 
     # Resume and modify configs at the earliest point.
     # The actual checkpoint path is in cfg.ckpt_path

From 6654272bc0c5017ea26dbe39ff7610d23c528bee Mon Sep 17 00:00:00 2001
From: Weilin Xu <weilin.xu@intel.com>
Date: Mon, 2 Jun 2025 16:46:44 -0700
Subject: [PATCH 15/18] Fix callback imports.

---
 mart/callbacks/__init__.py            | 1 +
 mart/callbacks/adversary_connector.py | 3 +--
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/mart/callbacks/__init__.py b/mart/callbacks/__init__.py
index 5e648370..d04a6d26 100644
--- a/mart/callbacks/__init__.py
+++ b/mart/callbacks/__init__.py
@@ -1,3 +1,4 @@
+# All Lightning callbacks dependent on lightning, so we don't import mart.callbacks by default.
 from ..utils.imports import _HAS_TORCHVISION
 from .adversary_connector import *
 from .eval_mode import *
diff --git a/mart/callbacks/adversary_connector.py b/mart/callbacks/adversary_connector.py
index d243da46..006fbe1f 100644
--- a/mart/callbacks/adversary_connector.py
+++ b/mart/callbacks/adversary_connector.py
@@ -12,8 +12,7 @@
 import torch
 from lightning.pytorch.callbacks import Callback
 
-from ..utils import MonkeyPatch
-from ..utils.pylogger import get_pylogger
+from ..utils import MonkeyPatch, get_pylogger
 
 logger = get_pylogger(__name__)
 

From 33cd7065d1da1578a45b0109425fbde79ccb9010 Mon Sep 17 00:00:00 2001
From: Weilin Xu <weilin.xu@intel.com>
Date: Mon, 2 Jun 2025 16:47:11 -0700
Subject: [PATCH 16/18] Conditional imports for lightnign utils.

---
 mart/utils/__init__.py | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/mart/utils/__init__.py b/mart/utils/__init__.py
index 986c6374..024d6e43 100644
--- a/mart/utils/__init__.py
+++ b/mart/utils/__init__.py
@@ -1,5 +1,11 @@
 # Only import components without external dependency.
 from .adapters import *
+from .imports import _HAS_LIGHTNING
 from .monkey_patch import *
 from .optimization import *
 from .silent import *
+
+if _HAS_LIGHTNING:
+    from .lightning import *
+    from .pylogger import *
+    from .rich_utils import *

From d48f017bf114fd20eb0edabcb5a5ffc91b7c8cf1 Mon Sep 17 00:00:00 2001
From: Weilin Xu <weilin.xu@intel.com>
Date: Mon, 2 Jun 2025 16:51:44 -0700
Subject: [PATCH 17/18] Move mart.utils.lightning.flatten_dict to mart.utils.

---
 mart/models/modular.py  |  2 +-
 mart/utils/__init__.py  |  1 +
 mart/utils/lightning.py | 20 --------------------
 mart/utils/utils.py     | 24 ++++++++++++++++++++++++
 tests/test_utils.py     |  2 +-
 5 files changed, 27 insertions(+), 22 deletions(-)
 create mode 100644 mart/utils/utils.py

diff --git a/mart/models/modular.py b/mart/models/modular.py
index ca10ab49..67cb9d5e 100644
--- a/mart/models/modular.py
+++ b/mart/models/modular.py
@@ -13,7 +13,7 @@
 
 from ..nn import SequentialDict
 from ..optim import OptimizerFactory
-from ..utils.lightning import flatten_dict
+from ..utils import flatten_dict
 from ..utils.optimization import configure_optimizers
 
 logger = logging.getLogger(__name__)
diff --git a/mart/utils/__init__.py b/mart/utils/__init__.py
index 024d6e43..6e05a092 100644
--- a/mart/utils/__init__.py
+++ b/mart/utils/__init__.py
@@ -4,6 +4,7 @@
 from .monkey_patch import *
 from .optimization import *
 from .silent import *
+from .utils import *
 
 if _HAS_LIGHTNING:
     from .lightning import *
diff --git a/mart/utils/lightning.py b/mart/utils/lightning.py
index f4a0a4ec..bc3ef69a 100644
--- a/mart/utils/lightning.py
+++ b/mart/utils/lightning.py
@@ -27,7 +27,6 @@
     "log_hyperparameters",
     "save_file",
     "task_wrapper",
-    "flatten_dict",
 ]
 
 log = pylogger.get_pylogger(__name__)
@@ -274,22 +273,3 @@ def get_resume_checkpoint(config: DictConfig) -> Tuple[DictConfig]:
         config = hydra.compose(config_name, overrides=overrides)
 
     return config
-
-
-def flatten_dict(d, delimiter="."):
-    def get_dottedpath_items(d: dict, parent: Optional[str] = None):
-        """Get pairs of the dotted path and the value from a nested dictionary."""
-        for name, value in d.items():
-            path = f"{parent}{delimiter}{name}" if parent else name
-            if isinstance(value, dict):
-                yield from get_dottedpath_items(value, parent=path)
-            else:
-                yield path, value
-
-    ret = {}
-    for key, value in get_dottedpath_items(d):
-        if key in ret:
-            raise KeyError(f"Key collision when flattening a dictionary: {key}")
-        ret[key] = value
-
-    return ret
diff --git a/mart/utils/utils.py b/mart/utils/utils.py
new file mode 100644
index 00000000..b19e87b0
--- /dev/null
+++ b/mart/utils/utils.py
@@ -0,0 +1,24 @@
+from typing import Optional
+
+__all__ = [
+    "flatten_dict",
+]
+
+
+def flatten_dict(d, delimiter="."):
+    def get_dottedpath_items(d: dict, parent: Optional[str] = None):
+        """Get pairs of the dotted path and the value from a nested dictionary."""
+        for name, value in d.items():
+            path = f"{parent}{delimiter}{name}" if parent else name
+            if isinstance(value, dict):
+                yield from get_dottedpath_items(value, parent=path)
+            else:
+                yield path, value
+
+    ret = {}
+    for key, value in get_dottedpath_items(d):
+        if key in ret:
+            raise KeyError(f"Key collision when flattening a dictionary: {key}")
+        ret[key] = value
+
+    return ret
diff --git a/tests/test_utils.py b/tests/test_utils.py
index b2aa3b44..cc1f6911 100644
--- a/tests/test_utils.py
+++ b/tests/test_utils.py
@@ -6,7 +6,7 @@
 
 import pytest
 
-from mart.utils.lightning import flatten_dict
+from mart.utils import flatten_dict
 
 
 def test_flatten_dict():

From 7fe6b027905b101ad733b047ad682b75ea8263ab Mon Sep 17 00:00:00 2001
From: Weilin Xu <weilin.xu@intel.com>
Date: Mon, 2 Jun 2025 16:59:31 -0700
Subject: [PATCH 18/18] Remove pylogger imports that depends on lightning.

---
 mart/attack/initializer/base.py   | 4 ----
 mart/attack/initializer/vision.py | 5 +++--
 2 files changed, 3 insertions(+), 6 deletions(-)

diff --git a/mart/attack/initializer/base.py b/mart/attack/initializer/base.py
index 49a5b173..9fd91091 100644
--- a/mart/attack/initializer/base.py
+++ b/mart/attack/initializer/base.py
@@ -10,10 +10,6 @@
 
 import torch
 
-from mart.utils import pylogger
-
-logger = pylogger.get_pylogger(__name__)
-
 
 class Initializer:
     """Initializer base class."""
diff --git a/mart/attack/initializer/vision.py b/mart/attack/initializer/vision.py
index 363141ef..0824fe26 100644
--- a/mart/attack/initializer/vision.py
+++ b/mart/attack/initializer/vision.py
@@ -4,14 +4,15 @@
 # SPDX-License-Identifier: BSD-3-Clause
 #
 
+import logging
+
 import torch
 import torchvision
 import torchvision.transforms.functional as F
 
-from ...utils import pylogger
 from .base import Initializer
 
-logger = pylogger.get_pylogger(__name__)
+logger = logging.getLogger(__name__)
 
 
 class Image(Initializer):