From d7b0764d4823d9a3380fc41a1de9913bf66cf9f8 Mon Sep 17 00:00:00 2001
From: Jeff Rasley <jerasley@microsoft.com>
Date: Thu, 12 Oct 2023 15:05:50 -0700
Subject: [PATCH] Update accelerator/cuda_accelerator.py

Co-authored-by: Michael Wyatt <michaelwyatt@microsoft.com>
---
 accelerator/cuda_accelerator.py | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/accelerator/cuda_accelerator.py b/accelerator/cuda_accelerator.py
index 0daf167e14e0..3ecbe9738f48 100644
--- 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):