tensorflow
import tensorflow as tf
print(tf.__version__)
print(tf.keras.__version__)
if tf.test.is_gpu_available(cuda_only=True):
device = "/gpu:0"
else:
device = "/cpu:0"
print(device)
Pytorch
import torch
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print(device)