diff --git a/.github/workflows/ci-testing.yml b/.github/workflows/ci-testing.yml
index 3924ba2..1520ab9 100644
--- a/.github/workflows/ci-testing.yml
+++ b/.github/workflows/ci-testing.yml
@@ -16,11 +16,15 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        os: ["ubuntu-22.04", "macOS-13", "windows-2022"]
+        os: ["ubuntu-24.04", "macOS-13", "windows-2022"]
         python-version: ["3.9", "3.12"]
         requires: ["latest"]
+        dependency: ["lightning"]
         include:
-          - { os: "ubuntu-20.04", python-version: "3.9", requires: "oldest" }
+          - { os: "ubuntu-22.04", python-version: "3.9", requires: "oldest", dependency: "lightning" }
+          - { os: "ubuntu-24.04", python-version: "3.10", requires: "latest", dependency: "pytorch_lightning" }
+          - { os: "ubuntu-24.04", python-version: "3.12", requires: "latest", dependency: "pytorch_lightning" }
+          - { os: "macOS-13", python-version: "3.12", requires: "latest", dependency: "pytorch_lightning" }
 
     # Timeout: https://stackoverflow.com/a/59076067/4521646
     timeout-minutes: 35
@@ -41,11 +45,23 @@ jobs:
           pip install 'lightning-utilities[cli]'
           python -m lightning_utilities.cli requirements set-oldest --req_files='["requirements.txt"]'
 
+      - name: Adjust requirements
+        run: |
+          pip install 'lightning-utilities[cli]' -U -q
+          python -m lightning_utilities.cli requirements replace-pkg \
+            --old_package="lightning" \
+            --new_package="${{matrix.dependency}}" \
+            --req_files='["_requirements/extra.txt"]'
+          cat _requirements/extra.txt
+
       - name: Install package & dependencies
         run: |
+          set -e
           pip --version
           pip install -e '.[test,extra]' -U -q --find-links $TORCH_URL
           pip list
+          # check that right package was installed
+          python -c "import ${{matrix.dependency}}; print(${{matrix.dependency}}.__version__)"
 
       - name: Tests with mocks
         run: |
diff --git a/src/litmodels/integrations/lightning_checkpoint.py b/src/litmodels/integrations/lightning_checkpoint.py
index f336c09..016f2d8 100644
--- a/src/litmodels/integrations/lightning_checkpoint.py
+++ b/src/litmodels/integrations/lightning_checkpoint.py
@@ -9,7 +9,8 @@
     from lightning.pytorch import Trainer
     from lightning.pytorch.callbacks import ModelCheckpoint
 elif _PYTORCHLIGHTNING_AVAILABLE:
-    from pytorch_lightning.callbacks import ModelCheckpoint, Trainer
+    from pytorch_lightning import Trainer
+    from pytorch_lightning.callbacks import ModelCheckpoint
 else:
     raise ModuleNotFoundError("No module named 'lightning' or 'pytorch_lightning'")