added requirement and skorch argument to avoid skorch unpickling errors in torch 2.6

Author: dholzmueller

docs/requirements.txt

@@ -27,7 +27,7 @@ seaborn>=0.0.13
 skorch>=0.15
 sphinx>=7.0
 sphinx_rtd_theme>=2.0
-torch>=2.0
+torch>=2.0,<2.6
 torchmetrics>=1.2.1
 tqdm
 tueplots>=0.0.12

pyproject.toml

@@ -26,7 +26,9 @@ classifiers = [
     "License :: OSI Approved :: Apache Software License",
 ]
 dependencies = [
-    "torch>=2.0",
+    # use <2.6 for now since it can run into pickling issues with skorch if the skorch version is too old
+    # see https://github.com/skorch-dev/skorch/commit/be93b7769d61aa22fb928d2e89e258c629bfeaf9
+    "torch>=2.0,<2.6",
     "numpy>=1.25,<2.0",
     "pandas>=2.0",
     "scikit-learn>=1.3,<1.6",

pytabkit/models/nn_models/rtdl_resnet.py

@@ -960,10 +960,8 @@ def create_regressor_skorch(
         model_class = RTDL_MLP
     else:
         raise ValueError(f'Model {model_name} not implemented here! Choose from "ft_transformer", "resnet", "mlp"')
-    model = nn_class(
-        model_class,
-        # Shuffle training data on each epoch
-        optimizer=optimizer,
+
+    new_kwargs = dict(optimizer=optimizer,
         batch_size=max(
             batch_size, 1
         ),  # if batch size is float, it will be reset during fit
@@ -974,8 +972,22 @@ def create_regressor_skorch(
         module__regression=True,
         module__categorical_indicator=None,  # will be change when fitted
         callbacks=callbacks,
-        **kwargs,
-    )
+        **kwargs)
+
+    try:
+        # try the torch_load_kwargs but it's only available in newer versions of skorch
+        model = nn_class(
+            model_class,
+            # Shuffle training data on each epoch
+            **new_kwargs,
+            torch_load_kwargs={'weights_only': False}, # quick-fix for pickling errors in torch>=2.6
+        )
+    except ValueError:
+        model = nn_class(
+            model_class,
+            # Shuffle training data on each epoch
+            **new_kwargs,
+        )
 
     return model