Python tf.test.is_built_with_cuda用法及代码示例

返回是否使用 CUDA (GPU) 支持构建 TensorFlow。

用法

tf.test.is_built_with_cuda()

此方法应仅用于使用 tf.test.TestCase 编写的测试。一个典型的用法是跳过应该只使用 CUDA (GPU) 运行的测试。

class MyTest(tf.test.TestCase):

  def test_add_on_gpu(self):
    if not tf.test.is_built_with_cuda():
      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 构建的。