Unverified Commit 2d61823f authored by Fanli Lin's avatar Fanli Lin Committed by GitHub
Browse files

[tests] add `require_torch_sdpa` for test that needs sdpa support (#30408)

* add cuda flag

* check for sdpa

* add bitsandbytes
parent 04ac3245
......@@ -374,6 +374,7 @@ class CohereModelTest(ModelTesterMixin, GenerationTesterMixin, PipelineTesterMix
@slow
class CohereIntegrationTest(unittest.TestCase):
@require_torch_multi_gpu
@require_bitsandbytes
def test_batched_4bit(self):
model_id = "CohereForAI/c4ai-command-r-v01-4bit"
......@@ -393,6 +394,7 @@ class CohereIntegrationTest(unittest.TestCase):
output = model.generate(**inputs, max_new_tokens=40, do_sample=False)
self.assertEqual(tokenizer.batch_decode(output, skip_special_tokens=True), EXPECTED_TEXT)
@require_torch_sdpa
def test_batched_small_model_logits(self):
# Since the model is very large, we created a random cohere model so that we can do a simple
# logits check on it.
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment