[test] Fix flaky test for tensorflow 2.7.0 on GPU (#1475)

Change-Id: If1a2a5cf58214f0cacee362e2806936ca6e271b3
deepjavalibrary · Jan 19, 2022 · 12a8c1c · 12a8c1c
1 parent 3b50f6b
commit 12a8c1c
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/...ration/src/main/java/ai/djl/integration/tests/ndarray/NDArrayElementArithmeticOpTest.java b/...ration/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 {