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.