Python tf.test.is_built_with_rocm用法及代碼示例

返回 TensorFlow 是否使用 ROCm (GPU) 支持構建。

用法

tf.test.is_built_with_rocm()

此方法應僅用於使用 tf.test.TestCase 編寫的測試。一個典型的用法是跳過應該隻使用 ROCm (GPU) 運行的測試。

class MyTest(tf.test.TestCase):

  def test_add_on_gpu(self):
    if not tf.test.is_built_with_rocm():
      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 官方二進製文件不是用 ROCm 構建的。