added mark for gpu in the rotary embedding test

mosaicml · Oct 20, 2023 · 5981ade · 5981ade
1 parent 0a9d3af
commit 5981ade
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/tests/test_rotary_embedding.py b/tests/test_rotary_embedding.py
@@ -13,6 +13,7 @@ def allclose_helper(t0: torch.Tensor,
     return torch.allclose(t0, t1, rtol=rtol, atol=atol)
 
 
+@pytest.mark.gpu
 @pytest.mark.parametrize('device', ['cpu', 'cuda'])
 @pytest.mark.parametrize('dtype', [torch.float32, torch.bfloat16])
 @pytest.mark.parametrize('rope_scaling_type',