Get number of gpus

gistfile1.txt

def gpu_count() -> int:
    """Get number of gpus."""
    import ctypes

    try:
        cuda = ctypes.CDLL("libcuda.so")
    except OSError:
        logging.exception("Loading libcuda")
        return 0

    CUDA_SUCCESS = 0
    CUDA_ERROR_NO_DEVICE = 100

    num_gpus = ctypes.c_int()
    result = ctypes.c_int()
    result = cuda.cuInit(0)
    if result != CUDA_SUCCESS:
        logging.error("cuInit returned an error: %d", result)
        return 0

    result = cuda.cuDeviceGetCount(ctypes.byref(num_gpus))
    if result != CUDA_SUCCESS:
        logging.error("cuDeviceGetCount returned an error: %d", result.value)
        return 0

    return num_gpus.value