test_cuda_tf2.py

#!/usr/bin/env python

import tensorflow as tf

from tensorflow.python.client import device_lib
print(str(device_lib.list_local_devices() ))

with tf.compat.v1.Graph().as_default():
        m1 = tf.constant([[1.0, 2.0]])
        m2 = tf.constant([[3.0], [4.0]])
        product = tf.matmul(m1, m2)

        with tf.compat.v1.Session() as session:
                result = session.run(product)

        print(result)