[Misc] Improve error message for incorrect pynvml (#12809)
Some checks failed
Create Release / Create Release (push) Has been cancelled
Some checks failed
Create Release / Create Release (push) Has been cancelled
Signed-off-by: youkaichao <youkaichao@gmail.com>
This commit is contained in:
@@ -41,7 +41,11 @@ def cuda_platform_plugin() -> Optional[str]:
|
|||||||
is_cuda = True
|
is_cuda = True
|
||||||
finally:
|
finally:
|
||||||
pynvml.nvmlShutdown()
|
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.
|
# CUDA is supported on Jetson, but NVML may not be.
|
||||||
import os
|
import os
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user