diff --git a/environment.yaml b/environment.yaml
index f02b80161..b44d4bb30 100644
--- a/environment.yaml
+++ b/environment.yaml
@@ -23,7 +23,7 @@ channels:
 dependencies:
   - pytorch=2.*
   - torchvision=0.*
-  - pytorch-lightning=2.*
+  - lightning=2.*
   - torchmetrics=0.*
   - hydra-core=1.*
   - rich=13.*
@@ -35,6 +35,7 @@ dependencies:
   # - neptune-client
   # - mlflow
   # - comet-ml
+  # - aim>=3.16.2 # no lower than 3.16.2, see https://github.com/aimhubio/aim/issues/2550
 
   - pip>=23
   - pip:
diff --git a/setup.py b/setup.py
index 51becc7af..f1826cc5b 100644
--- a/setup.py
+++ b/setup.py
@@ -9,7 +9,7 @@
     author="",
     author_email="",
     url="https://github.com/user/project",
-    install_requires=["pytorch-lightning", "hydra-core"],
+    install_requires=["lightning", "hydra-core"],
     packages=find_packages(),
     # use this to customize global commands available in the terminal after installing the package
     entry_points={