Fix mypy

Author: adamjstewart

.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)/

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)

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(),