Skip to content

cuda

util.cuda

init_cuda_when_available 

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.

Source code in SaigeToolkit/util/cuda.py 
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
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 as e:
        # if any error occurs while checking cuda.is_available, raise CudaDeviceError
        raise CudaDeviceError from e

    if _is_available:
        try:
            # init cuda
            torch.cuda.init()
        except Exception as exc:
            # if any error occurs while init cuda, raise CudaDeviceError
            raise CudaDeviceError from exc