Test if different DL frameworks work well with your current CUDA installation

test_pytorch.py

#!/usr/bin/env python3
"""Tests PyTorch CUDA capabilities"""

import sys

import torch


def main():
    """Main function"""
    print("A", sys.version)
    print("B", torch.__version__)
    print("C", torch.cuda.is_available())
    print("D", torch.backends.cudnn.enabled)
    device = torch.device("cuda")
    print("E", torch.cuda.get_device_properties(device))
    print("F", torch.tensor([1.0, 2.0]).cuda())


if __name__ == "__main__":
    main()

test_tensorflow.py

#!/usr/bin/env python3
"""Module doc"""

import sys
import tensorflow as tf
import numpy as np

def main():
    """Main function"""
    print(tf.config.list_physical_devices("GPU"))
    print("A", sys.version)
    print("B", tf.__version__)
    model = tf.keras.applications.VGG16(weights='imagenet')
    res = model.predict(np.random.rand(2, 224, 224, 3))
    print("C", res.shape)

if __name__ == "__main__":
    main()