cuda-check.py

import ctypes
import ctypes.util

_libcuda = ctypes.CDLL(ctypes.util.find_library('cudart'))

if __name__ == '__main__':
    # You should put the path to the CUDA toolkit library into LD_LIBRARY_PATH envvar
    # or into the system-wide configuration at /etc/ld.so.conf.d/.
    count = ctypes.c_int(0)

    # This will initialize a CUDA device context,
    # which holds the handle to the device until this program exits.
    ret = _libcuda.cudaGetDeviceCount(ctypes.byref(count))
    assert ret == 0, 'cudaGetDeviceCount() has failed (error code {}).'.format(ret)
    print('Number of GPUs detected: {}'.format(count.value))