[Hardware][NVIDIA] Add non-NVML CUDA mode for Jetson (#9735)

Signed-off-by: Conroy Cheers <conroy@corncheese.org>
This commit is contained in:
Conroy Cheers
2024-11-27 05:26:28 +11:00
committed by GitHub
parent db66e018ea
commit f5792c7c4a
3 changed files with 159 additions and 91 deletions

View File

@@ -28,7 +28,15 @@ try:
finally:
pynvml.nvmlShutdown()
except Exception:
pass
# CUDA is supported on Jetson, but NVML may not be.
import os
def cuda_is_jetson() -> bool:
return os.path.isfile("/etc/nv_tegra_release") \
or os.path.exists("/sys/class/tegra-firmware")
if cuda_is_jetson():
is_cuda = True
is_rocm = False