Merge pull request #8 from Alexei-V-Ivanov-AMD/main_20240517

Main 20240517
Alexei-V-Ivanov-AMD · May 18, 2024 · 1422420 · 1422420
2 parents 2e32cf7 + 68b68c6
commit 1422420
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/tests/spec_decode/e2e/conftest.py b/tests/spec_decode/e2e/conftest.py
@@ -6,8 +6,11 @@
 import pytest
 import ray
 import torch
-from pynvml import (nvmlDeviceGetHandleByIndex, nvmlDeviceGetMemoryInfo,
-                    nvmlInit)
+from vllm.utils import is_hip
+
+if (not is_hip()):
+    from pynvml import (nvmlDeviceGetHandleByIndex, nvmlDeviceGetMemoryInfo,
+                        nvmlInit)
 
 from vllm import LLM
 from vllm.engine.arg_utils import AsyncEngineArgs