diff --git a/tests/compile/test_dynamic_shapes_compilation.py b/tests/compile/test_dynamic_shapes_compilation.py index 96c3f49aba..7f725c14f2 100644 --- a/tests/compile/test_dynamic_shapes_compilation.py +++ b/tests/compile/test_dynamic_shapes_compilation.py @@ -9,6 +9,7 @@ import pytest import torch from tests.models.utils import check_logprobs_close +from tests.utils import wait_for_rocm_memory_to_settle from vllm import LLM, SamplingParams from vllm.compilation.decorators import support_torch_compile from vllm.config import CompilationConfig, VllmConfig, set_current_vllm_config @@ -104,6 +105,7 @@ def test_dynamic_shapes_compilation( gc.collect() torch.accelerator.empty_cache() torch.accelerator.synchronize() + wait_for_rocm_memory_to_settle() eager_model = LLM(model=model_name, enforce_eager=True, max_model_len=1024) eager_outputs = []