Use gpu for tf computation when available

gpu_availability_test.py

import unittest
import tensorflow as tf
import sys
import os

class TestGPU(unittest.TestCase):
    def setUp(self):
        os.environ["CUDA_VISIBLE_DEVICES"]='0'
        
    def test_gpu(self): 
        if not tf.test.is_built_with_cuda():
            print('cuda not supported')

        if not tf.test.is_gpu_available():
            print('gpu not available')
        else:
            # Creates a graph.
            with tf.device('/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)
            # Creates a session with log_device_placement set to True.
            sess = tf.Session(config=tf.ConfigProto(log_device_placement=True))
            # Runs the op.
            print(sess.run(c))


if __name__ == 'main':
    unittest.main()