Update accelerator/cuda_accelerator.py

Co-authored-by: Michael Wyatt <[email protected]>
microsoft · Oct 12, 2023 · d7b0764 · d7b0764
1 parent 26c2bcd
commit d7b0764
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/accelerator/cuda_accelerator.py b/accelerator/cuda_accelerator.py
@@ -30,12 +30,13 @@ def __init__(self):
     def _init_pynvml(self):
         global pynvml
         try:
-            import pynvml
+            import pynvml as tmp_pynvml
         except ImportError:
             return
         try:
-            pynvml.nvmlInit()
-        except pynvml.NVMLError:
+            tmp_pynvml.nvmlInit()
+            pynvml = tmp_pynvml
+        except tmp_pynvml.NVMLError:
             return
 
     def is_synchronized_device(self):