make tests pass mps (#528)

TransformerLensOrg · Mar 26, 2024 · 93c2246 · 93c2246
1 parent 6673d88
commit 93c2246
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/tests/acceptance/test_hooked_transformer.py b/tests/acceptance/test_hooked_transformer.py
@@ -311,8 +311,12 @@ def check_dtype(dtype, margin, no_processing=False):
         gc.collect()
 
 
+@pytest.mark.skipif(
+    torch.backends.mps.is_available() or not torch.cuda.is_available(),
+    reason="some operations unsupported by MPS: https://github.com/pytorch/pytorch/issues/77754 or no GPU",
+)
 @pytest.mark.parametrize("dtype", [torch.float64, torch.float32])
-def test_dtypes(dtype):
+def test_dtype_float(dtype):
     check_dtype(dtype, margin=5e-4)