Racapef/torchvision

.gitlab-ci.yml

@@ -109,7 +109,13 @@ test:
   before_script:
     - *default-before-script
     - |
-      if [ "$PYTORCH_IMAGE" == "1.13.1-cuda11.6-cudnn8-devel" ]; then
+      if [ "${TORCH_OVERRIDE:-}" == "2.0.0" ]; then
+        # No torch 2.0.0 + CUDA 11.8 image is published: reuse the 2.2.0 cu118
+        # image (Python 3.10, CUDA 11.8) and install torch 2.0.0+cu118.
+        CUDA_VERSION_NAME=cu118
+        TORCH_VERSION=2.0.0+cu118
+        TORCHVISION_VERSION=0.15.1+cu118
+      elif [ "$PYTORCH_IMAGE" == "1.13.1-cuda11.6-cudnn8-devel" ]; then
         CUDA_VERSION_NAME=cu116
         TORCH_VERSION=1.13.1+cu116
         TORCHVISION_VERSION=0.14.1+cu116
@@ -160,7 +166,7 @@ test:
       update_dict(base, delta)
       Path("pyproject.toml").write_text(tomlkit.dumps(base))
       EOF
-    - uv sync --group=test --reinstall
+    - uv sync --group=test --extra=attn --reinstall
     - *check-torch-cuda
     - |
       PYTEST_ARGS=(--cov=compressai --capture=no tests)
@@ -174,6 +180,9 @@ test:
       - PYTORCH_IMAGE:
           - "1.13.1-cuda11.6-cudnn8-devel"
           - "2.2.0-cuda11.8-cudnn8-devel"
+      # torch 2.0.0+cu118 / torchvision 0.15.1 on the 2.2.0 cu118 image.
+      - PYTORCH_IMAGE: "2.2.0-cuda11.8-cudnn8-devel"
+        TORCH_OVERRIDE: "2.0.0"
   tags:
     - docker
 

pyproject.toml

@@ -44,7 +44,11 @@ dependencies = [
   "torch>=1.13.1",
   "torch>=2.6.0; python_version >= '3.12'",
   "torchvision>=0.14.1,<0.20; python_version < '3.9'",
-  "torchvision>=0.17.0; python_version >= '3.9'",
+  # torchvision 0.14.1-0.16.x are supported below Python 3.12 (exercised by the
+  # torch 1.13.1 and 2.0.0 CI jobs). torchvision 0.17 is the first release with
+  # Python 3.12 wheels.
+  "torchvision>=0.14.1; python_version >= '3.9' and python_version < '3.12'",
+  "torchvision>=0.17.0; python_version >= '3.12'",
   "tqdm",
   "typing-extensions>=4.0.0",
   "wheel>=0.32.0", # For --no-build-isolation.

tests/test_mlicv2_layers.py

@@ -37,15 +37,20 @@
 
 import compressai.layers as layers
 
-from compressai.models._helpers.mlicv2 import (
-    ContextReweighting,
-    GSCModule,
-    HGCPModule,
-    RoPE2D,
-    SimpleTokenMixing,
-    STMAnalysis,
-    STMSynthesis,
-)
+try:
+    from compressai.models._helpers.mlicv2 import (
+        ContextReweighting,
+        GSCModule,
+        HGCPModule,
+        RoPE2D,
+        SimpleTokenMixing,
+        STMAnalysis,
+        STMSynthesis,
+    )
+except ModuleNotFoundError as err:
+    if err.name != "timm":
+        raise
+    pytestmark = pytest.mark.skip(reason="MLICv2 layers require the [attn] extra")
 
 
 def _round_trip_output(