fix another

mosaicml · Nov 18, 2023 · 21031b5 · 21031b5
1 parent b64166b
commit 21031b5
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/tests/test_model.py b/tests/test_model.py
@@ -1352,11 +1352,11 @@ def test_generate_with_past_kv(attn_impl: str, pos_emb_config: dict,
                                                              hf_config.d_model)
 
 
-@pytest.mark.parametrize('attn_impl,device', [
+@pytest.mark.parametrize('attn_impl', [
     'torch',
-    pytest.param('flash', 'gpu', marks=pytest.mark.gpu),
-    pytest.param('triton', 'gpu', marks=pytest.mark.gpu),
-    pytest.param('torch', 'gpu', marks=pytest.mark.gpu),
+    pytest.param('flash', marks=pytest.mark.gpu),
+    pytest.param('triton', marks=pytest.mark.gpu),
+    pytest.param('torch', marks=pytest.mark.gpu),
 ])
 @pytest.mark.parametrize('generation_kwargs', [{
     'max_new_tokens': 2,