View source on GitHub |
Returns whether TensorFlow was built with ROCm (GPU) support.
tf.test.is_built_with_rocm()
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 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 official binary is NOT built with ROCm.