make explicit gpu dep

huggingface · Jan 13, 2025 · e2cb0b9 · e2cb0b9
1 parent b8c34d9
commit e2cb0b9
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/tests/models/bamba/test_modeling_bamba.py b/tests/models/bamba/test_modeling_bamba.py
@@ -23,6 +23,7 @@
 from transformers import AutoTokenizer, BambaConfig, is_torch_available
 from transformers.testing_utils import (
     require_torch,
+    require_torch_gpu,
     slow,
     torch_device,
 )
@@ -493,6 +494,7 @@ def _prepare_model_kwargs(input_ids, attention_mask, signature):
 
 @slow
 @require_torch
+@require_torch_gpu
 class BambaModelIntegrationTest(unittest.TestCase):
     model = None
     tokenizer = None