m

tensorflow · Dec 25, 2023 · 0f7ffb5 · 0f7ffb5
1 parent c3fac82
commit 0f7ffb5
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/...flow_recommenders_addons/dynamic_embedding/python/kernel_tests/horovod_sync_train_test.py b/...flow_recommenders_addons/dynamic_embedding/python/kernel_tests/horovod_sync_train_test.py
@@ -65,6 +65,9 @@ def get_emb_sequential_model(emb_t, opt, *args, **kwargs):
   return model
 
 
+is_gpu_available = test_util.is_gpu_available()
+
+
 class HorovodTest(test.TestCase):
 
   @test_util.deprecated_graph_mode_only
@@ -90,6 +93,8 @@ def test_all_to_all_embedding_trainable(self):
     if (tf.__version__ == "2.11.0" or tf.__version__ == "2.11.1"):
       self.skipTest(
           "The save function doesn't work with TF 2.11, skip the test.")
+    if not is_gpu_available:
+      self.skipTest('Only test when gpu is available.')
     if (is_macos() and is_arm64()):
       self.skipTest(
           "Apple silicon devices don't support synchronous training based on Horovod."