disable partial attention tests on rocm

facebookresearch · Mar 1, 2024 · 6752f07 · 6752f07
1 parent b585563
commit 6752f07
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/tests/test_mem_eff_attention.py b/tests/test_mem_eff_attention.py
@@ -2688,6 +2688,7 @@ def paged_attention_run_inner(
     torch.testing.assert_close(y_swapped, y_packed)
 
 
+@disable_on_rocm
 @sm80_or_better_only
 @pytest.mark.parametrize(
     "op",
@@ -2748,7 +2749,7 @@ def test_merge_attentions_nobias(
         assert lse is None
 
 
-@sm80_or_better_only
+@disable_on_rocm
 @sm80_or_better_only
 @pytest.mark.parametrize(
     "dtype,op",