[bugfix] fix full graph tests (vllm-project#10581)

Signed-off-by: youkaichao <[email protected]>
neuralmagic · Nov 22, 2024 · db100c5 · db100c5
1 parent 11fcf0e
commit db100c5
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/tests/compile/utils.py b/tests/compile/utils.py
@@ -4,7 +4,7 @@
 
 from tests.quantization.utils import is_quant_method_supported
 from vllm import LLM, SamplingParams
-from vllm.config import CompilationConfig, CompilationLevel
+from vllm.config import CompilationLevel
 from vllm.platforms import current_platform
 
 TEST_MODELS = [
@@ -85,7 +85,7 @@ def check_full_graph_support(model,
               enforce_eager=True,
               tensor_parallel_size=tp_size,
               disable_custom_all_reduce=True,
-              compilation_config=CompilationConfig(level=optimization_level),
+              compilation_config=optimization_level,
               **model_kwargs)
 
     outputs = llm.generate(prompts, sampling_params)