Internal change

PiperOrigin-RevId: 405455038

Internal change
PiperOrigin-RevId: 405455038
26d8e705 · Abdullah Rashwan · A. Unique TensorFlower · 6cd536dc · 26d8e705
Commit 26d8e705 authored Oct 25, 2021 by Abdullah Rashwan Committed by A. Unique TensorFlower Oct 25, 2021
Hide whitespace changes
Inline Side-by-side

Showing with 4 additions and 2 deletions

official/vision/beta/serving/export_module_factory_test.py official/vision/beta/serving/export_module_factory_test.py +4 -2

No files found.
--- a/official/vision/beta/serving/export_module_factory_test.py
+++ b/official/vision/beta/serving/export_module_factory_test.py
@@ -107,8 +107,10 @@ class ImageClassificationExportTest(tf.test.TestCase, parameterized.TestCase):
    # The imported model should contain any trackable attrs that the original
    # model had.
    self.assertTrue(hasattr(imported.model, 'test_trackable'))
-    self.assertAllClose(out['logits'].numpy(), expected_logits.numpy())
-    self.assertAllClose(out['probs'].numpy(), expected_prob.numpy())
+    self.assertAllClose(
+        out['logits'].numpy(), expected_logits.numpy(), rtol=1e-04, atol=1e-04)
+    self.assertAllClose(
+        out['probs'].numpy(), expected_prob.numpy(), rtol=1e-04, atol=1e-04)


 if __name__ == '__main__':