Check Google Colab GPU Type for RAPIDS Compatibility

check_colab_gpu.py

import pynvml

pynvml.nvmlInit()
handle = pynvml.nvmlDeviceGetHandleByIndex(0)
device_name = pynvml.nvmlDeviceGetName(handle)

if (device_name != b'Tesla T4') and (device_name != b'Tesla P4') and (device_name != b'Tesla P100-PCIE-16GB'):
  raise Exception("""
                     Unfortunately this instance does not have a T4, P4 or P100 GPU.\n
                     Please make sure you've configured Colab to request a GPU instance type.\n
                     Sometimes Colab allocates a Tesla K80 instead of a T4, P4 or P100. Resetting the instance.If you get a K80 GPU, try Runtime -> Reset all runtimes...""")
else:
  print('Woo! You got the right kind of GPU!')