diff --git a/integration/src/main/java/ai/djl/integration/tests/training/listener/EarlyStoppingListenerTest.java b/integration/src/main/java/ai/djl/integration/tests/training/listener/EarlyStoppingListenerTest.java
index eefe7121025..9aee2661411 100644
--- a/integration/src/main/java/ai/djl/integration/tests/training/listener/EarlyStoppingListenerTest.java
+++ b/integration/src/main/java/ai/djl/integration/tests/training/listener/EarlyStoppingListenerTest.java
@@ -178,7 +178,7 @@ public void testEarlyStoppingStopsOnEpoch1AsMaxDurationIs1ms() throws Exception
                     // epoch
                     EasyTrain.fit(trainer, 10, trainMnistDataset, testMnistDataset);
                 } catch (EarlyStoppingListener.EarlyStoppedException e) {
-                    Assert.assertTrue(e.getMessage().contains("ms elapsed >= 1 maxMillis"));
+                    Assert.assertTrue(e.getMessage().contains("ms elapsed >="));
                     Assert.assertTrue(e.getStopEpoch() < 10); // Stop epoch is before 10
                 }