fix import path

Signed-off-by: Jonathan Nitisastro <jonathancn@anyscale.com>
ray-project · Nov 13, 2023 · e7e2d67 · e7e2d67
1 parent 39c410e
commit e7e2d67
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/python/ray/tests/accelerators/test_nvidia_gpu.py b/python/ray/tests/accelerators/test_nvidia_gpu.py
@@ -3,7 +3,11 @@
 import pytest
 
 from ray._private.accelerators import NvidiaGPUAcceleratorManager
-from mock_pynvml import DeviceHandleMock, PyNVMLMock, patch_mock_pynvml
+from ray.tests.accelerators.mock_pynvml import (
+    DeviceHandleMock,
+    PyNVMLMock,
+    patch_mock_pynvml,
+)
 
 GPU_MOCK_DATA = [
     DeviceHandleMock(