From 12a8c1c0419ed6082ee6f70320d3126d4cbfd5c3 Mon Sep 17 00:00:00 2001
From: Frank Liu <frankfliu2000@gmail.com>
Date: Wed, 19 Jan 2022 13:50:08 -0800
Subject: [PATCH] [test] Fix flaky test for tensorflow 2.7.0 on GPU (#1475)

Change-Id: If1a2a5cf58214f0cacee362e2806936ca6e271b3
---
 .../tests/ndarray/NDArrayElementArithmeticOpTest.java         | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/integration/src/main/java/ai/djl/integration/tests/ndarray/NDArrayElementArithmeticOpTest.java b/integration/src/main/java/ai/djl/integration/tests/ndarray/NDArrayElementArithmeticOpTest.java
index 43dbb72f831..d52a1a2154d 100644
--- a/integration/src/main/java/ai/djl/integration/tests/ndarray/NDArrayElementArithmeticOpTest.java
+++ b/integration/src/main/java/ai/djl/integration/tests/ndarray/NDArrayElementArithmeticOpTest.java
@@ -55,7 +55,7 @@ private void testCornerCase(
             array1 = manager.create(new float[] {x1, x1});
             array2 = manager.create(x2);
             NDArray expected = manager.create(new float[] {y, y});
-            Assert.assertEquals(arrayArithmeticOp.apply(array1, array2), expected);
+            Assertions.assertAlmostEquals(arrayArithmeticOp.apply(array1, array2), expected);
 
             // test zero-dim with zero-dim
             array1 = manager.create(new Shape(4, 0, 1));
@@ -119,7 +119,7 @@ private void testReverseScalarCornerCase(
         float x2 = 3f;
         NDArray ndArray = manager.create(x1);
         NDArray result = arrayArithmeticOp.apply(x2, ndArray);
-        Assert.assertEquals(result.getFloat(), arithmeticOp.apply(x2, x1).floatValue());
+        Assertions.assertAlmostEquals(result.getFloat(), arithmeticOp.apply(x2, x1));
         if (inplace) {
             Assert.assertSame(ndArray, result);
         } else {