diff --git a/api/src/test/kotlin/org/jetbrains/kotlinx/dl/api/core/integration/InferenceModelTest.kt b/api/src/test/kotlin/org/jetbrains/kotlinx/dl/api/core/integration/InferenceModelTest.kt
index 946139be7..ddeeb1fac 100644
--- a/api/src/test/kotlin/org/jetbrains/kotlinx/dl/api/core/integration/InferenceModelTest.kt
+++ b/api/src/test/kotlin/org/jetbrains/kotlinx/dl/api/core/integration/InferenceModelTest.kt
@@ -220,9 +220,18 @@ class InferenceModelTest {
                     it.predict(train.getX(0))
                 }
             assertEquals(
-                "Reshape functions is missed!",
+                "Reshape functions is missed! Define and set up the reshape function to transform initial data to the model input.",
                 exception.message
             )
+
+            val exception2 =
+                Assertions.assertThrows(IllegalArgumentException::class.java) {
+                    it.predictSoftly(train.getX(0))
+                }
+            assertEquals(
+                "Reshape functions is missed! Define and set up the reshape function to transform initial data to the model input.",
+                exception2.message
+            )
         }
     }