[BUG] Ensure compatibility with lightning > 2.6 and above (#625)

Fixes #623, Fixes #621 

Adds a wrapper around `load_from_checkpoint` to ensure compatibility
with all versions

Author: phoeenniixx

.github/workflows/testing.yml

@@ -53,6 +53,41 @@ jobs:
       - name: Run test-suite
         run: python -m pytest -v
 
+  test-deps-2024:
+    runs-on: ${{ matrix.os }}
+    strategy:
+      fail-fast: false
+      matrix:
+        os: [ubuntu-latest]
+        python-version: ["3.12"]
+
+    steps:
+      - uses: actions/checkout@v6
+
+      - run: git remote set-branches origin 'main'
+
+      - run: git fetch --depth 1
+
+      - name: Install uv
+        uses: astral-sh/setup-uv@v7
+        with:
+          enable-cache: true
+
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v6
+        with:
+          python-version: ${{ matrix.python-version }}
+
+      - name: Install main package & dependencies
+        run: uv pip install -e .[dev,dependencies_2024]
+        env:
+          UV_SYSTEM_PYTHON: 1
+
+      - name: Show installed packages
+        run: uv pip list
+
+      - name: Run test-suite
+        run: python -m pytest -v
 
   test-all-extras:
     runs-on: ${{ matrix.os }}

pyproject.toml

@@ -43,7 +43,7 @@ dependencies = [
     "numpy<=3.0.0",
     "pandas>=1.1.5,<3.0.0",
     "scikit-learn>=1.3.0,<2.0",
-    "pytorch-lightning>=2.0.0,<2.5.0",
+    "pytorch-lightning>=2.0.0,<2.7.0",
     "scipy>=1.8,<2.0",
     "omegaconf>=2.3.0",
     "torchmetrics>=0.10.0,<1.9.0",
@@ -80,6 +80,19 @@ notebooks = [
     "matplotlib>3.1",
 ]
 
+# Core Dep set in Nov 2024
+dependencies_2024 = [
+    "torch==2.5.0",
+    "numpy==2.2.0",
+    "pandas==2.2.3",
+    "scikit-learn==1.6.0",
+    "pytorch-lightning==2.4.0",
+    "scipy==1.14.1",
+    "omegaconf==2.3.0",
+    "torchmetrics==1.5.2",
+    "einops==0.8.0",
+]
+
 
 [project.urls]
 Homepage = "https://github.com/pytorch-tabular/pytorch_tabular"

src/pytorch_tabular/models/base_model.py

@@ -8,6 +8,7 @@
 from abc import ABCMeta, abstractmethod
 from functools import partial
 from typing import Any, Callable, Dict, List, Optional, Tuple, Type, Union
+from pathlib import Path
 
 import numpy as np
 import pytorch_lightning as pl
@@ -206,6 +207,28 @@ def embedding_layer(self):
     def head(self):
         raise NotImplementedError("head property needs to be implemented by inheriting classes")
 
+    @classmethod
+    def load_from_checkpoint(
+        cls,
+        checkpoint_path: Union[str, Path],
+        map_location=None,
+        strict=True,
+        **kwargs,
+    ):
+        from skbase.utils.dependencies import _check_soft_dependencies
+
+        if not _check_soft_dependencies("pytorch_lightning<2.6", severity="none"):
+            if "weights_only" not in kwargs:
+                kwargs["weights_only"] = False
+        else:
+            kwargs.pop("weights_only", None)
+        return super().load_from_checkpoint(
+            checkpoint_path,
+            map_location=map_location,
+            strict=strict,
+            **kwargs,
+        )
+
     def _check_and_verify(self):
         assert hasattr(self, "backbone"), "Model has no attribute called `backbone`"
         assert hasattr(self.backbone, "output_dim"), "Backbone needs to have attribute `output_dim`"

src/pytorch_tabular/ssl_models/base_model.py

@@ -5,7 +5,8 @@
 
 import warnings
 from abc import ABCMeta, abstractmethod
-from typing import Dict, Optional
+from typing import Dict, Optional, Union
+from pathlib import Path
 
 import pytorch_lightning as pl
 import torch
@@ -151,6 +152,28 @@ def forward(self, x: Dict):
     def featurize(self, x: Dict):
         pass
 
+    @classmethod
+    def load_from_checkpoint(
+        cls,
+        checkpoint_path: Union[str, Path],
+        map_location=None,
+        strict=True,
+        **kwargs,
+    ):
+        from skbase.utils.dependencies import _check_soft_dependencies
+
+        if not _check_soft_dependencies("pytorch_lightning<2.6", severity="none"):
+            if "weights_only" not in kwargs:
+                kwargs["weights_only"] = False
+        else:
+            kwargs.pop("weights_only", None)
+        return super().load_from_checkpoint(
+            checkpoint_path,
+            map_location=map_location,
+            strict=strict,
+            **kwargs,
+        )
+
     def predict(self, x: Dict, ret_model_output: bool = True):  # ret_model_output only for compatibility
         assert ret_model_output, "ret_model_output must be True in case of SSL predict"
         return self.featurize(x)