SUMIN/HuggingFace_transformer
1
0
Fork 0
Code Issues Pull Requests Actions 7 Packages Projects Releases Wiki Activity

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

Browse Source 
fix cuda-only tests
This commit is contained in:
Fanli Lin
2024-12-18 17:14:22 +08:00
committed by GitHub
parent 1eee1cedfd
commit c7e48053aa
1 changed files with 3 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
tests/models/rag/test_modeling_rag.py
View File

@@ -33,7 +33,7 @@ from transformers.testing_utils import (
require_sentencepiece, require_sentencepiece,
require_tokenizers, require_tokenizers,
require_torch, require_torch,
require_torch_non_multi_gpu, require_torch_non_multi_accelerator,
slow, slow,
torch_device, torch_device,
) )
@@ -678,7 +678,7 @@ class RagDPRT5Test(RagTestMixin, unittest.TestCase):
@require_retrieval @require_retrieval
@require_sentencepiece @require_sentencepiece
@require_tokenizers @require_tokenizers
@require_torch_non_multi_gpu @require_torch_non_multi_accelerator
class RagModelIntegrationTests(unittest.TestCase): class RagModelIntegrationTests(unittest.TestCase):
def tearDown(self): def tearDown(self):
super().tearDown() super().tearDown()
@@ -1002,7 +1002,7 @@ class RagModelIntegrationTests(unittest.TestCase):
torch_device torch_device
) )
if torch_device == "cuda": if torch_device != "cpu":
rag_token.half() rag_token.half()
input_dict = tokenizer( input_dict = tokenizer(