Unverified Commit 0408efc6 authored by youkaichao's avatar youkaichao Committed by GitHub
Browse files

[Misc] Improve error message for incorrect pynvml (#12809)


Signed-off-by: default avataryoukaichao <youkaichao@gmail.com>
parent 449d1bce
......@@ -41,7 +41,11 @@ def cuda_platform_plugin() -> Optional[str]:
is_cuda = True
finally:
pynvml.nvmlShutdown()
except Exception:
except Exception as e:
if "nvml" not in e.__class__.__name__.lower():
# If the error is not related to NVML, re-raise it.
raise e
# CUDA is supported on Jetson, but NVML may not be.
import os
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment