tf.test.is_built_with_gpu_support

Returns whether TensorFlow was built with GPU (CUDA or ROCm) support.

View aliases

Compat aliases for migration

See Migration guide for more details.

tf.compat.v1.test.is_built_with_gpu_support

tf.test.is_built_with_gpu_support()

This method should only be used in tests written with tf.test.TestCase. A typical usage is to skip tests that should only run with GPU.

class MyTest(tf.test.TestCase):

  def test_add_on_gpu(self):
    if not tf.test.is_built_with_gpu_support():
      self.skipTest("test is only applicable on GPU")

    with tf.device("GPU:0"):
      self.assertEqual(tf.math.add(1.0, 2.0), 3.0)

TensorFlow official binary is built with CUDA GPU support.