Update: PyTorch v2.6.0 (#1093)

- Add torch 2.6.0 to CI
- Remove torch 2.2.2
- Update torch install instructions, as they no longer provide conda
  packages
- Add test for new default of weights_only
- Update pickle file test artifact (explained in #1092)
- Update some comments
- Conditionally install triton 3.1 for torch < 2.6

Author: BenjaminBossan

Diff for: .github/workflows/testing.yml

@@ -21,7 +21,7 @@ jobs:
       fail-fast: false  # don't cancel all jobs when one fails
       matrix:
         python_version: ['3.9', '3.10', '3.11', '3.12']
-        torch_version: ['2.2.2+cpu', '2.3.1+cpu', '2.4.1+cpu', '2.5.1+cpu']
+        torch_version: ['2.3.1+cpu', '2.4.1+cpu', '2.5.1+cpu', '2.6.0+cpu']
         os: [ubuntu-latest]
 
     steps:
@@ -32,13 +32,18 @@ jobs:
         python-version: ${{ matrix.python_version }}
     - name: Install dependencies
       # TODO remove numpy version constraint once we no longer support PyTorch < 2.3
+      # TODO remove triton 3.1 install if torch < 2.6 no longer supported, see #1093
       run: |
         python -m pip install --upgrade pip
         python -m pip install -r requirements-dev.txt
         python -m pip install -r requirements.txt
         python -m pip install --force-reinstall -U "numpy<2.0.0"
         python -m pip install pytest-pretty
         python -m pip install torch==${{ matrix.torch_version }} -f https://download.pytorch.org/whl/torch
+        TORCH_VERSION_MAJOR_MINOR=$(python -c "import torch; v=torch.__version__.split('+')[0]; print('.'.join(v.split('.')[:2]))")
+        if [[ $(echo "$TORCH_VERSION_MAJOR_MINOR < 2.6" | bc -l) -eq 1 ]]; then
+          python -m pip install "triton==3.1"
+        fi
         python -m pip list
     - name: Install skorch
       run: |

Diff for: README.rst

@@ -177,9 +177,9 @@ To install skorch from source using conda, proceed as follows:
 
     git clone https://github.com/skorch-dev/skorch.git
     cd skorch
-    conda create -n skorch-env python=3.10
+    conda create -n skorch-env python=3.12
     conda activate skorch-env
-    conda install -c pytorch pytorch
+    python -m pip install torch
     python -m pip install -r requirements.txt
     python -m pip install .
 
@@ -189,9 +189,9 @@ If you want to help developing, run:
 
     git clone https://github.com/skorch-dev/skorch.git
     cd skorch
-    conda create -n skorch-env python=3.10
+    conda create -n skorch-env python=3.12
     conda activate skorch-env
-    conda install -c pytorch pytorch
+    python -m pip install torch
     python -m pip install -r requirements.txt
     python -m pip install -r requirements-dev.txt
     python -m pip install -e .
@@ -239,10 +239,10 @@ instructions for PyTorch, visit the `PyTorch website
 <http://pytorch.org/>`__. skorch officially supports the last four
 minor PyTorch versions, which currently are:
 
-- 2.2.2
 - 2.3.1
 - 2.4.1
 - 2.5.1
+- 2.6.0
 
 However, that doesn't mean that older versions don't work, just that
 they aren't tested. Since skorch mostly relies on the stable part of
@@ -252,9 +252,6 @@ In general, running this to install PyTorch should work:
 
 .. code:: bash
 
-    # using conda:
-    conda install pytorch pytorch-cuda -c pytorch
-    # using pip
     python -m pip install torch
 
 ==================

Diff for: docs/user/installation.rst

@@ -98,10 +98,10 @@ instructions for PyTorch, visit the `PyTorch website
 <http://pytorch.org/>`__. skorch officially supports the last four
 minor PyTorch versions, which currently are:
 
-- 2.2.2
 - 2.3.1
 - 2.4.1
 - 2.5.1
+- 2.6.0
 
 However, that doesn't mean that older versions don't work, just that
 they aren't tested. Since skorch mostly relies on the stable part of
@@ -111,7 +111,4 @@ In general, running this to install PyTorch should work:
 
 .. code:: bash
 
-    # using conda:
-    conda install pytorch pytorch-cuda -c pytorch
-    # using pip
     python -m pip install torch

Diff for: skorch/tests/net_cuda.pkl

@@ -3075,9 +3075,8 @@ def test_torch_load_kwargs_auto_weights_false_pytorch_lt_2_6(
     ):
         # Same test as
         # test_torch_load_kwargs_auto_weights_only_false_when_load_params but
-        # without monkeypatching get_default_torch_load_kwargs. There is no
-        # corresponding test for >= 2.6.0 since it's not clear yet if the switch
-        # will be made in that version.
+        # without monkeypatching get_default_torch_load_kwargs. The default is
+        # weights_only=False.
         # See discussion in 1063.
         from skorch._version import Version
 
@@ -3098,6 +3097,32 @@ def test_torch_load_kwargs_auto_weights_false_pytorch_lt_2_6(
         del call_kwargs['map_location']  # we're not interested in that
         assert call_kwargs == expected_kwargs
 
+    def test_torch_load_kwargs_auto_weights_true_pytorch_ge_2_6(
+            self, net_cls, module_cls, monkeypatch, tmp_path
+    ):
+        # Same test as
+        # test_torch_load_kwargs_auto_weights_false_pytorch_lt_2_6 but
+        # with weights_only=True, since it's the new default
+        # See discussion in 1063.
+        from skorch._version import Version
+
+        # TODO remove once torch 2.5.0 is no longer supported
+        if Version(torch.__version__) < Version('2.6.0'):
+            pytest.skip("Test only for torch >= 2.6.0")
+
+        net = net_cls(module_cls).initialize()
+        net.save_params(f_params=tmp_path / 'params.pkl')
+        state_dict = net.module_.state_dict()
+        expected_kwargs = {"weights_only": True}
+
+        mock_torch_load = Mock(return_value=state_dict)
+        monkeypatch.setattr(torch, "load", mock_torch_load)
+        net.load_params(f_params=tmp_path / 'params.pkl')
+
+        call_kwargs = mock_torch_load.call_args_list[0].kwargs
+        del call_kwargs['map_location']  # we're not interested in that
+        assert call_kwargs == expected_kwargs
+
     def test_torch_load_kwargs_forwarded_to_torch_load_unpickle(
             self, net_cls, module_cls, monkeypatch, tmp_path
     ):

Diff for: skorch/tests/test_net.py

@@ -775,10 +775,10 @@ def get_default_torch_load_kwargs():
     """Returns the kwargs passed to torch.load that correspond to the current
     torch version.
 
-    The plan is to switch from weights_only=False to True in PyTorch version
-    2.6.0, but depending on what happens, this may require updating.
+    PyTorch switches from weights_only=False to True in version 2.6.0.
 
     """
+    # TODO: Remove once PyTorch 2.5 is no longer supported
     version_torch = Version(torch.__version__)
     version_default_switch = Version('2.6.0')
     if version_torch >= version_default_switch: