use callback to terminate minibatch tests

SciML · Sep 14, 2024 · c012905 · c012905
1 parent 2a803ff
commit c012905
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/test/minibatch.jl b/test/minibatch.jl
@@ -21,7 +21,7 @@ end
 
 function callback(state, l) #callback function to observe training
     display(l)
-    return false
+    return l < 1e-3
 end
 
 u0 = Float32[200.0]
@@ -100,7 +100,7 @@ function callback(st, l, pred; doplot = false)
         scatter!(pl, t, pred[1, :], label = "prediction")
         display(plot(pl))
     end
-    return false
+    return l < 1e-3
 end
 
 optfun = OptimizationFunction(loss_adjoint,