Fixing redundant import pynvml.

Alexei-V-Ivanov-AMD · May 18, 2024 · 68b68c6 · 68b68c6
1 parent e2351a4
commit 68b68c6
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