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

Signed-off-by: youkaichao <youkaichao@gmail.com>
This commit is contained in:
youkaichao
2025-02-06 15:23:50 +08:00
committed by GitHub
parent 449d1bce02
commit 0408efc6d0

View File

@ -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