move decode into try block

rjzamora · Dec 11, 2024 · 9d72fe7 · 9d72fe7
1 parent 6eabc90
commit 9d72fe7
Showing 1 changed file with 5 additions and 3 deletions.
diff --git a/distributed/diagnostics/nvml.py b/distributed/diagnostics/nvml.py
@@ -104,9 +104,11 @@ def init_once():
             NVML_STATE = NVMLState.DISABLED_LIBRARY_NOT_FOUND
             return
 
-        if _in_wsl() and parse_version(
-            pynvml.nvmlSystemGetDriverVersion().decode()
-        ) < parse_version(MINIMUM_WSL_VERSION):
+        try:
+            driver_vsn = pynvml.nvmlSystemGetDriverVersion().decode()
+        except AttributeError:
+            driver_vsn = pynvml.nvmlSystemGetDriverVersion()
+        if _in_wsl() and parse_version(driver_vsn) < parse_version(MINIMUM_WSL_VERSION):
             NVML_STATE = NVMLState.DISABLED_WSL_INSUFFICIENT_DRIVER
             return
         else: