From 2d61823fa2f5bf18c56a9c1a0730e9d16b5d395d Mon Sep 17 00:00:00 2001 From: Fanli Lin Date: Tue, 23 Apr 2024 17:39:38 +0800 Subject: [PATCH] [tests] add `require_torch_sdpa` for test that needs sdpa support (#30408) * add cuda flag * check for sdpa * add bitsandbytes --- tests/models/cohere/test_modeling_cohere.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/models/cohere/test_modeling_cohere.py b/tests/models/cohere/test_modeling_cohere.py index 3e86ffe9d96..2165943e796 100644 --- a/tests/models/cohere/test_modeling_cohere.py +++ b/tests/models/cohere/test_modeling_cohere.py @@ -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.