diff --git a/tests/e2e/integrations/test_llm_compressor.py b/tests/e2e/integrations/test_llm_compressor.py index f4317559d..f43d4a938 100644 --- a/tests/e2e/integrations/test_llm_compressor.py +++ b/tests/e2e/integrations/test_llm_compressor.py @@ -30,12 +30,12 @@ MODELS = [ @pytest.mark.parametrize( "save_compressed", [True, False], ids=["save_compressed", "save_uncompressed"] ) -@require_llmcompressor class TestLLMCompressorIntegration: """ e2e tests for axolotl.integrations.llm_compressor.LLMCompressorPlugin """ + @require_llmcompressor @require_torch_2_4_1 def test_llmcompressor_plugin( self, temp_dir, base_model: str, save_compressed: bool diff --git a/tests/e2e/utils.py b/tests/e2e/utils.py index 0228957c1..61df1d8fe 100644 --- a/tests/e2e/utils.py +++ b/tests/e2e/utils.py @@ -105,7 +105,7 @@ def require_vllm(test_case): return False return unittest.skipUnless( - is_vllm_installed(), "test requires a vllm to be installed" + is_vllm_installed(), "test requires vllm to be installed" )(test_case) @@ -123,7 +123,7 @@ def require_llmcompressor(test_case): return False return unittest.skipUnless( - is_llmcompressor_installed(), "test requires a llmcompressor to be installed" + is_llmcompressor_installed(), "test requires llmcompressor to be installed" )(test_case)