Skip to content

Instantly share code, notes, and snippets.

@supercrazysam
Created September 7, 2020 01:09
Show Gist options
  • Save supercrazysam/03cdb18ffc181d37e791cc71486d6219 to your computer and use it in GitHub Desktop.
Save supercrazysam/03cdb18ffc181d37e791cc71486d6219 to your computer and use it in GitHub Desktop.
CUDA test script
import tensorflow as tf
if tf.test.gpu_device_name():
print('Default GPU Device: {}'.format(tf.test.gpu_device_name()))
else:
print("Please install GPU version of TF")
import tensorflow as tf
with tf.device('/gpu:0'):
a = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[2, 3], name='a')
b = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[3, 2], name='b')
c = tf.matmul(a, b)
with tf.Session() as sess:
print (sess.run(c))
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment