Testing GPU Setup

Use these code samples to test that your GPU setup works with several common deep learning libraries. The specific versions of libraries depend on the particular GPU used and the GPU driver version. You can use this testing for GPU setup using ML Runtimes or Legacy Engines.

  1. Go to a project that is using the CUDA engine and click Open Workbench.
  2. Launch a new session with GPUs.
  3. Run the following command in the workbench command prompt to verify that the driver was installed correctly:
    ! /usr/bin/nvidia-smi
  4. Use any of the following code samples to confirm that the new engine works with common deep learning libraries.


!pip3 install torch==1.4.0
from torch import cuda
assert cuda.is_available()
assert cuda.device_count() > 0


!pip3 install tensorflow-gpu==2.1.0
from tensorflow.python.client import device_lib
assert 'GPU' in str(device_lib.list_local_devices())


!pip3 install keras
from keras import backend
assert len(backend.tensorflow_backend._get_available_gpus()) > 0