Bring back mps tests

weigertlab · Jun 17, 2024 · 8f38dbd · 8f38dbd
1 parent b35b831
commit 8f38dbd
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tests/test_pretrained.py b/tests/test_pretrained.py
@@ -9,7 +9,7 @@
 
 
 @pytest.mark.parametrize("name", ["ctc", "general_2d"])
-@pytest.mark.parametrize("device", ["cpu", "cuda"])
+@pytest.mark.parametrize("device", ["cpu", "cuda", "mps"])
 def test_pretrained(name, device):
     """Each pretrained model should run on all (available) device."""
     # TODO mps support