init cuda when available.
if cuda device has error and requires reboot, it will raise CudaDeviceError.
Raises:
Source code in SaigeToolkit/util/cuda.py
| def init_cuda_when_available():
"""init cuda when available.
if cuda device has error and requires reboot, it will raise CudaDeviceError.
Raises:
CudaDeviceError: cuda device error.
"""
try:
# check cuda device is available
_is_available = torch.cuda.is_available()
except Exception:
# if any error occurs while checking cuda.is_available, raise CudaDeviceError
raise CudaDeviceError
if _is_available:
try:
# init cuda
torch.cuda.init()
except Exception:
# if any error occurs while init cuda, raise CudaDeviceError
raise CudaDeviceError
|