tf.test.is_built_with_gpu_support

Возвращает значение, указывающее, было ли TensorFlow скомпилировано с поддержкой GPU (CUDA или ROCm).

tf.test.is_built_with_gpu_support()

Этот метод следует использовать только в тестах, написанных с помощью tf.test.TestCase. Типичное использование — пропуск тестов, которые должны выполняться только с 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 скомпилирована с поддержкой графических процессоров CUDA.