diff --git a/tests/models/language/pooling/conftest.py b/tests/models/language/pooling/conftest.py index 2da4c65133..ceb21b5a40 100644 --- a/tests/models/language/pooling/conftest.py +++ b/tests/models/language/pooling/conftest.py @@ -33,6 +33,8 @@ def release_gpu_memory_between_tests(): gc.collect() if torch.accelerator.is_available(): torch.accelerator.empty_cache() + if current_platform.is_rocm(): + return try: wait_for_memory_to_settle() except ValueError as e: diff --git a/tests/models/language/pooling_mteb_test/conftest.py b/tests/models/language/pooling_mteb_test/conftest.py index 5a31a5f415..ce34d82563 100644 --- a/tests/models/language/pooling_mteb_test/conftest.py +++ b/tests/models/language/pooling_mteb_test/conftest.py @@ -9,6 +9,7 @@ import torch from tests.utils import wait_for_memory_to_settle from vllm.logger import init_logger +from vllm.platforms import current_platform logger = init_logger(__name__) @@ -26,6 +27,8 @@ def release_gpu_memory_between_tests(): gc.collect() if torch.accelerator.is_available(): torch.accelerator.empty_cache() + if current_platform.is_rocm(): + return try: wait_for_memory_to_settle() except ValueError as e: diff --git a/tests/models/multimodal/pooling/conftest.py b/tests/models/multimodal/pooling/conftest.py index 32f598b30a..fdd58e2fba 100644 --- a/tests/models/multimodal/pooling/conftest.py +++ b/tests/models/multimodal/pooling/conftest.py @@ -9,6 +9,7 @@ import torch from tests.utils import wait_for_memory_to_settle from vllm.logger import init_logger +from vllm.platforms import current_platform logger = init_logger(__name__) @@ -26,6 +27,8 @@ def release_gpu_memory_between_tests(): gc.collect() if torch.accelerator.is_available(): torch.accelerator.empty_cache() + if current_platform.is_rocm(): + return try: wait_for_memory_to_settle() except ValueError as e: