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 (#35607)

Browse Source 
* intial commit

* remove unrelated files

* further remove

* Update test_trainer.py

* fix style
This commit is contained in:
Fanli Lin
2025-01-13 21:48:39 +08:00
committed by GitHub
parent e6f9b03464
commit 2fa876d2d8
18 changed files with 57 additions and 47 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
tests/models/mixtral/test_modeling_mixtral.py
View File

@@ -22,6 +22,7 @@ from transformers import MixtralConfig, is_torch_available
from transformers.testing_utils import (
require_flash_attn,
require_torch,
require_torch_accelerator,
require_torch_gpu,
slow,
torch_device,
@@ -471,7 +472,7 @@ class MixtralIntegrationTest(unittest.TestCase):
cls.cuda_compute_capability_major_version = torch.cuda.get_device_capability()[0]
@slow
@require_torch_gpu
@require_torch_accelerator
def test_small_model_logits(self):
model_id = "hf-internal-testing/Mixtral-tiny"
dummy_input = torch.LongTensor([[0, 1, 0], [0, 1, 0]]).to(torch_device)
@@ -507,7 +508,7 @@ class MixtralIntegrationTest(unittest.TestCase):
)
@slow
@require_torch_gpu
@require_torch_accelerator
def test_small_model_logits_batched(self):
model_id = "hf-internal-testing/Mixtral-tiny"
dummy_input = torch.LongTensor([[0, 0, 0, 0, 0, 0, 1, 2, 3], [1, 1, 2, 3, 4, 5, 6, 7, 8]]).to(torch_device)