Check CUDA/cuDNN info

check-cuda.py

# Run as:
# wget -O check-cuda.py https://gist.github.com/fauxneticien/343b1dd7b68a30cb6f8983dacac28721/raw && python check-cuda.py

# Adapted from https://amytabb.com/til/2020/10/05/cudnn-pytorch/

import torch

print("Is cuda available?", torch.cuda.is_available())

print("CUDA version:", torch.version.cuda)

print("cuDNN version:", torch.backends.cudnn.version())

print("cuDNN enabled? ", torch.backends.cudnn.enabled)

print("Try creating a tensor 'x' and moving to GPU...")
x = torch.rand(5, 3).to('cuda')

print("What device is 'x' on?", x.device)