remove .numpy which failed in gpu/distributed env

tensorflow · Apr 30, 2024 · 185c34d · 185c34d
1 parent 697341f
commit 185c34d
Showing 1 changed file with 3 additions and 8 deletions.
diff --git a/...ow_recommenders_addons/dynamic_embedding/python/kernel_tests/ragged_embedding_ops_test.py b/...ow_recommenders_addons/dynamic_embedding/python/kernel_tests/ragged_embedding_ops_test.py
@@ -17,11 +17,6 @@ def test_fill_empty_rows(self):
     expected_filled = tf.ragged.constant([[1, 2, 3], [0], [4], [0], [5, 6]],
                                          dtype=tf.int32)
     expected_empty = tf.constant([False, True, False, True, False])
-
-    self.assertTrue(
-        tf.reduce_all(filled_ragged_tensor.to_tensor() ==
-                      expected_filled.to_tensor()).numpy(),
-        "Filled tensors do not match")
-    self.assertTrue(
-        tf.reduce_all(is_row_empty == expected_empty).numpy(),
-        "Empty row flags do not match")
+    tf.debugging.assert_equal(filled_ragged_tensor.to_tensor(),
+                              expected_filled.to_tensor())
+    tf.debugging.assert_equal(is_row_empty, expected_empty)