diff --git a/test/diffeqfluxtests.jl b/test/diffeqfluxtests.jl
index 692a1f382..5af34ed14 100644
--- a/test/diffeqfluxtests.jl
+++ b/test/diffeqfluxtests.jl
@@ -84,7 +84,7 @@ function loss_neuralode(p)
 end
 
 iter = 0
-callback = function (st, l)
+callback = function (st, l, pred)
     global iter
     iter += 1
 
diff --git a/test/minibatch.jl b/test/minibatch.jl
index f818f4ee1..5a4c1af01 100644
--- 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-2
 end
 
 u0 = Float32[200.0]
@@ -58,11 +58,11 @@ optfun = OptimizationFunction(loss_adjoint,
     Optimization.AutoZygote())
 optprob = OptimizationProblem(optfun, pp, train_loader)
 
-res1 = Optimization.solve(optprob,
-    Optimization.Sophia(; η = 0.5,
-        λ = 0.0), callback = callback,
-    maxiters = 1000)
-@test 10res1.objective < l1
+# res1 = Optimization.solve(optprob,
+#     Optimization.Sophia(; η = 0.5,
+#         λ = 0.0), callback = callback,
+#     maxiters = 1000)
+# @test 10res1.objective < l1
 
 optfun = OptimizationFunction(loss_adjoint,
     Optimization.AutoForwardDiff())
@@ -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,