[tests] make cuda-only tests device-agnostic (#35222)

fix cuda-only tests
huggingface · Dec 18, 2024 · c7e4805 · c7e4805
1 parent 1eee1ce
commit c7e4805
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/tests/models/rag/test_modeling_rag.py b/tests/models/rag/test_modeling_rag.py
@@ -33,7 +33,7 @@
     require_sentencepiece,
     require_tokenizers,
     require_torch,
-    require_torch_non_multi_gpu,
+    require_torch_non_multi_accelerator,
     slow,
     torch_device,
 )
@@ -678,7 +678,7 @@ def config_and_inputs(self):
 @require_retrieval
 @require_sentencepiece
 @require_tokenizers
-@require_torch_non_multi_gpu
+@require_torch_non_multi_accelerator
 class RagModelIntegrationTests(unittest.TestCase):
     def tearDown(self):
         super().tearDown()
@@ -1002,7 +1002,7 @@ def test_rag_token_generate_batch(self):
             torch_device
         )
 
-        if torch_device == "cuda":
+        if torch_device != "cpu":
             rag_token.half()
 
         input_dict = tokenizer(