chore: use cpu wheel for torchvision

.github/workflows/tests.yml

@@ -30,6 +30,8 @@ jobs:
             poetry source add torch-cpu --priority=explicit --no-interaction https://download.pytorch.org/whl/cpu
             torch_version=$(poetry show torch --no-ansi | grep 'version' | awk '{print $3}')
             poetry add --lock --no-interaction --source torch-cpu "torch@$torch_version"
+            torchvision_version=$(poetry show torchvision --no-ansi | grep 'version' | awk '{print $3}')
+            poetry add --lock --no-interaction --source torch-cpu "torchvision@$torchvision_version"
       - name: Load cached venv
         id: cached-poetry-dependencies
         uses: actions/cache@v4

pyproject.toml

@@ -19,11 +19,11 @@ migrate-ckpt = {git = "https://github.com/bdvllrs/migrate-ckpt.git", rev = "v0.2
 cfg-tools = {git = "https://github.com/bdvllrs/cfg-tools.git", rev = "v0.3.0"}
 pillow = "^10.3.0"
 numpy = "^1.26"
-torch = ">=2.0.1,<2.5.0"
+torch = ">=2.0.1,<2.6.0"
 wandb = "^0.18.3"
 lightning = ">=2.1.0"
 pydantic = "^2.6.0"
-torchvision = ">=0.15.2,<0.20"
+torchvision = ">=0.15.2,<0.21"
 rich = "^13.4.2"
 
 [tool.poetry.group.dev.dependencies]