relax unit test to 1e-05

Author: FDecaYed

distributed_embeddings/python/layers/dist_model_parallel_test.py

@@ -189,7 +189,7 @@ def run_and_test(self, ref_model, ref_inputs, test_model, test_inputs):
 
     for ref_w, test_w in zip(ref_weights, test_weights):
       # assert close here since order of accumulations(inputs and batch dim) might have changed
-      self.assertAllClose(tf.convert_to_tensor(ref_w), tf.convert_to_tensor(test_w))
+      self.assertAllClose(tf.convert_to_tensor(ref_w), tf.convert_to_tensor(test_w), 1e-05, 1e-05)
 
   def test_broadcast(self):
     tf.keras.utils.set_random_seed(int(time.time()) + self.hvd_rank)

tests/dist_model_parallel_test.py

@@ -189,7 +189,7 @@ def run_and_test(self, ref_model, ref_inputs, test_model, test_inputs):
 
     for ref_w, test_w in zip(ref_weights, test_weights):
       # assert close here since order of accumulations(inputs and batch dim) might have changed
-      self.assertAllClose(tf.convert_to_tensor(ref_w), tf.convert_to_tensor(test_w))
+      self.assertAllClose(tf.convert_to_tensor(ref_w), tf.convert_to_tensor(test_w), 1e-05, 1e-05)
 
   def test_broadcast(self):
     tf.keras.utils.set_random_seed(int(time.time()) + self.hvd_rank)