Check if using tensorflow gpu version

tf.test.is_built_with_cuda()

Note: not sure how this funciton behaves with rocm