From b259ebf9b0c16cf28452fefe3eefdb265fd65dc7 Mon Sep 17 00:00:00 2001
From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com>
Date: Thu, 30 Jan 2025 10:00:15 +0000
Subject: [PATCH] Bump the torch group in /requirements with 2 updates (#2546)

* Bump the torch group in /requirements with 2 updates

Bumps the torch group in /requirements with 2 updates: [torch](https://github.com/pytorch/pytorch) and [torchvision](https://github.com/pytorch/vision).


Updates `torch` from 2.5.1 to 2.6.0
- [Release notes](https://github.com/pytorch/pytorch/releases)
- [Changelog](https://github.com/pytorch/pytorch/blob/main/RELEASE.md)
- [Commits](https://github.com/pytorch/pytorch/compare/v2.5.1...v2.6.0)

Updates `torchvision` from 0.20.1 to 0.21.0
- [Release notes](https://github.com/pytorch/vision/releases)
- [Commits](https://github.com/pytorch/vision/compare/v0.20.1...v0.21.0)

---
updated-dependencies:
- dependency-name: torch
  dependency-type: direct:production
  update-type: version-update:semver-minor
  dependency-group: torch
- dependency-name: torchvision
  dependency-type: direct:production
  update-type: version-update:semver-minor
  dependency-group: torch
...

Signed-off-by: dependabot[bot] <support@github.com>

* Fix mypy

---------

Signed-off-by: dependabot[bot] <support@github.com>
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
Co-authored-by: Adam J. Stewart <ajstewart426@gmail.com>
---
 .pre-commit-config.yaml    | 2 +-
 requirements/required.txt  | 4 ++--
 tests/conftest.py          | 2 +-
 tests/models/test_croma.py | 2 +-
 4 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index bc948edcb59..e1bd1a15d67 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -29,7 +29,7 @@ repos:
           - pillow>=10.4.0
           - pytest>=6.1.2
           - scikit-image>=0.22.0
-          - torch>=2.3
+          - torch>=2.6
           - torchmetrics>=0.10
           - torchvision>=0.18
         exclude: (build|data|dist|logo|logs|output)/
diff --git a/requirements/required.txt b/requirements/required.txt
index b00e5df8987..4cdb4896f2f 100644
--- a/requirements/required.txt
+++ b/requirements/required.txt
@@ -17,6 +17,6 @@ rtree==1.3.0
 segmentation-models-pytorch==0.4.0
 shapely==2.0.6
 timm==1.0.14
-torch==2.5.1
+torch==2.6.0
 torchmetrics==1.6.1
-torchvision==0.20.1
+torchvision==0.21.0
diff --git a/tests/conftest.py b/tests/conftest.py
index d55a972ced1..3f117bdb2ff 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -27,4 +27,4 @@ def matplotlib_backend() -> None:
 
 @pytest.fixture(autouse=True)
 def torch_hub(tmp_path: Path) -> None:
-    torch.hub.set_dir(tmp_path)  # type: ignore[no-untyped-call]
+    torch.hub.set_dir(tmp_path)
diff --git a/tests/models/test_croma.py b/tests/models/test_croma.py
index 83f63770339..ba857a57fbf 100644
--- a/tests/models/test_croma.py
+++ b/tests/models/test_croma.py
@@ -18,7 +18,7 @@
 )
 
 
-def save_model(model: torch.nn.Module, path: Path) -> None:
+def save_model(model: CROMA, path: Path) -> None:
     state_dict = {
         's1_encoder': model.s1_encoder.state_dict(),
         's1_GAP_FFN': model.s1_GAP_FFN.state_dict(),