fix: Remove cuda from test_emb

gustaveroussy · Aug 22, 2024 · 687e6c0 · 687e6c0
1 parent a57b13c
commit 687e6c0
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tests/test_emb.py b/tests/test_emb.py
@@ -81,7 +81,7 @@ def setup_class(self) -> None:
     @pytest.mark.parametrize("coord_dir,arch_name,custom_model,pretrained_weights,transforms_dict,batch_size,num_workers,device", [
         (None,"clam", None, "IMAGENET1K_V2", {"totensor": {}, "normalize": {"mean": [0.485, 0.456, 0.406], "std": [0.229, 0.224, 0.225]}}, 1, 1, "cpu"),
         ("./coords", "phikon", None, None, None, 8, 2, "cpu"),
-        ("./coords", None, resnet18(pretrained=True), None, None, 4, 4, "cuda")])
+        ("./coords", None, resnet18(pretrained=True), None, None, 4, 4, "cpu")])
     def test_extract_model_based_embeddings(
         self,
         coord_dir,