diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml
index ff287e15bd..7a4e4387ae 100644
--- a/.buildkite/pipeline.yml
+++ b/.buildkite/pipeline.yml
@@ -55,7 +55,7 @@ steps:
       queue: "juliagpu"
       cuda: "*"
     if: build.message !~ /\[skip docs\]/ && !build.pull_request.draft
-    timeout_in_minutes: 1000
+    timeout_in_minutes: 1440
 
 env:
   SECRET_DOCUMENTER_KEY: "yBJHrf5zUPu3VFotb0W4TRdauMTGgNUei3ax0xrVTmcrrP3EX8zSGaj9MNZji9H6JqyNEhCqZMGcHhxR5XK0f97YjhAp5rDNpwV6FbIXY8FXgFyIOLAenYUklta1W6fNM7KTz3Dq3UnKBAprKhQBwUw3SldTuvtm+IhiVT2rFmADqxMSQcv+5LivYEgAFrrd6LX+PHhonj38VN46z5Bi3TOIGAnByVlZssX7cOwaRg/8TloLPsWAUlQiPr89Vdlow9k6SQV8W9mf00/Rq4LFd1Eq1EYTCSmReVawMxVpXh1mj7MSabf9ppVWYOmjP5Rzatk8YFWlJf80HVvzY7tXRQ==;U2FsdGVkX1/UmKrEQkyqoZY6GckLGmpV5SsPa2byZhZbomiJKIrLntTIHK9GclgfCJ1MDCKhuo3hMVxq+NKVm0QmnMZk0Hlzszx8Up5pbcuSb2BA0Wm7IJzZ5/2uXfdBOXIejwk+GUabJtSpiDlYJUW/8bOWYZUwl9rmrOM44tPesEcs5yMZPShIowZnJqrvQUcWXZ/+aZjftX+Pe7nP0KddzYRYzhIJVzYmU394V9MWqZDqaNU19TJqnL8dNQxFLfyDibXuCqSD9QjV2I/iCfwr1xI81h11vmaatpZCDUUGYyxLRh1w5BuT1hHvLqktoAGliBD2lSOEB5KbwJu0nQCbEXPsM6CpMb4JygFloViQJ1L0kb91JEZIrnP6ruM1rrgrRu8Ufb+hob+BsCiEhSxwLy6ynOll2ZggMU/BVjQjkVQ3TmxBLvJ4T3QwzCSPpyN6lyHgSWNodN2Fn+PaWd0Sc0gSs8BY9PmOOc9m5ErVCrsscTr7aqcIbLZBnNA3y+mtLN5Vpuk3bmv9Bb9SAllXMLFFd604wtlxJi2Ecrb+6b4rc+QUmr5gnYqOYLCCOruFJfvMY63eFuuHWHKT+qroiGeuUhprUJDZUzhKUKnFFhhUZKtI11RAhVikZMPyMAsW4+gshQkEjXTRZqBpadpMh+c0uO8V2tRZAPIn7GXSdsCaGWbzL6yVZx79mM0W4sDKAp0Xs2wc04FgPbDu2sHGA+VxokrGctRovGVhSELx65aAj7x/ByhYdIByPCxHa7TNBUHb7n/4XLw8KIzVKr6jX2ep5m3KlYdjI7uxq8Hlpeu0hCRG3tdCqwOZsEOm3yhC3B/jsrKLzOsDP/x3ByAp8RvSVPP9WsWP55CxZUvc+q5LiVXBc7PhUC4nRB5/FMykjm6PboB92Y0DkP8Wql09FDr3vs8B3TkVLzOvzcz888oZTQpTaoixaAlVBs51al4N7UXhp5BInUCUIkknIyAEzXgK/5SpzixVoEkeNPkrMqg4hDaSHlKu0VDuqcP0Uv/9IKf/qs0+XK+2v9QBgqV16upbHK2EptgII3QJpXf2sq/IQTPXq3aly3NnpPUcIZJZtfG/Wk7qOFg/NhBVMvXWkrTRwQhJ5VXFTP0kXVpbgBml4Eq/zw+tAn5mmtieVaeFonZgqCIa+gi+tWiMy2V3aTubEYUGWTb3WOtxMt9i3bu9KhvOIr+fwCgpYUG9Vb/6v7H84zt5HT59sNlo9J8Acih8EfVZseC5JVF6ugxfnHc8BBOtvuUUFtOjIWwOgcfCiPXvsZdMQh0Ow3u9pYyJUZ3s+enHkBwtsu3+kXBkeL77eggslREVUz2g9eK8G5sKwSCsltgv8HQbHYARkXqq14Unb3NNakvab7BrQ2znWy7zU4P04Thtqn2fOqiAOUxuGRF5iNnkSnIZavtataT6rezB1efn4V2HYANcR4JoGj1JBXWa/dOJAYVnRLB7pbpS94yjbEMjRMB5aWSnjf6FzaDnXwfUAjAHvHNZsLxVhtIejyPQhg2kbSftwWWqw9zVvc3dz2a18V+MyNakcRiRS0CEqSl/L8vIhTppGultqzJZCKJKMAIzUE5+Q1EeDYL1kEImtSZ3hVXq4YaSB4q/bNKbDLG4pWs7IO6DFR88j5KIPfMPy15jgP6v+w0QZh8JNPQLveXvmYU4T9vxalT8i1P/4Nc2tI1A8HUv0hVoNu6V0ugTxukMLJLLFrZ80Ic7xCxNzqlzzcMhsYOHhcR4fZCdDtmoNfZm066hAkRcQJ0yNbiv7GUxrybGzer+N+s7QtS7/AGxuVW1MNQlsqdwTL2KTOvkZWHYB5cHpfqeS6zSPczeEiOIgT1fRXv3kYRBaJ4Dk7aWwXuCBbN2vAhRCEjcJC6QXE4BUukycNacytP1+HhCBeouxdPli9jrPIz2YH0iy7kc47XPiJr7zR1KAza3M3boau6xeb/why5zV7gHzB08qAxhm+pBLm4ERdKSwe/KAdGX5M0hiqMLHceUwJgzWEgukJfntxeZnP1rFJnTEAbkBy/CKtEmdr4IJYIFZU59IE9WOeYgrzl677JoGblkJ2B1sj1U8DbsRHIN+muVdAmYu+PBft0Zxih4JNe/rgmC/hNpDClMEPIEk4euRLf3xl1OHCOcWfEKuhwV/wIwJ0dtQzjN97g6a9IbSlupLAwPDZM925hC7qYicuzrF0Nj64GeOSMb4WIiEGpgH8TWOYkgxea+RoNLu0MEErcz26jqnV1QS8YGEFtT8k0lnhivg+SXIcgdVMoyavFVjqP4FeNm0aL0vE5+odOzEIzKKVNvHqofO4HbrRmlbAorS9OfzRlHbzJznIWO+8yyQc6oKyeT92THh4+wFYXQzkw0bvHsJZ08OymCAnIP+dZCzOSJ/MzcI3BjmcMZcHasPS6CfgSRtm7o8GJvyZljfW/z4zdy6+HzGLxZcgiwDc4qODYBEMdSRdDrxzpENZ4/IK3JTTcafsrRgbqi1nYyadQQx5A5xFhyYZq04iaDxx8GmDV6v4l4Hp/iPHoG0mp3yjxrt5hUjEyLW/+5lZXCnBxIfEiV/vdZBXGJxZs3oiATTOjMQxQYbbcyEs02WnFqRsMxDwlTsUnhbnEPUb9vkQhJjkHAsOt2V7csH7ojwlJPCp9baWV6iqSbvtQu5sSwWItOJflxiw2NCVvvMhGjOQOb8sFf6yeAKHOi+2vk0T7Kkr5JziPZ1L2PDmiZHwMhjVwd2uIZP6pBuQtoyxxn6040nUX5QwHjVO7RamVmqOLoKJFQHYWlRs1FHSxK6BQWH/1MeSDvCBWfiFnm6wWUMWr9NKlSOMFPWEnVdQ+eu83yoSRVT0U7WNoSw/tyK1KB64DL6Z7VffWB1AvnMZ1uvuSFtkEjHOTrcmDkGzSDQs61qO8kmDiTqlpMDVbv7DmhiXpBAC2agwLw/xh1m3xoRNetTNxowMVRjokZCrevoPLCcxrRDl0tZz9g/fF2q9rMRIAEhwWTYC+WgX4KQ4Xj4BpFbx1d3G2oHEAIJItyIFHHoKzwKJl+InNJEdXZUCEpdE3IiI2agnbP9W/1sSRWKE1Ub0KukhK7GYBIwdnwx0GgqqLYTRNrA8PnKLkSQp4ri9XJRSxI52jqbMP/S3x2ogIbvyYeQXbCwS7jloEMSgDSAQcTPZpPEzR5tHZG/XMMYHZX/h+aARdsaQJq7UNoIAJ8zrwkWnjNKSqrpF3Wfn/uOCYHXsyHepa/4f9cf0mtklGa4nSZPV8nVa0jvXzt2lzmg7uur0/ysa8Wl2LAZpkcTLlZ1tbFrbdBbcibPGi4r0QYJ6BM60yjfarwq/WnFHY2BLIpJKOxxA/7ko6kZ05t+fe/IqZnkxHqH9PsdpCN63H+py0S3tMOijrULzRMNjalF1ywPNm2Ugl/EiBbx7wv2G30wMdk1pdKbjgiGdq2wF1nPb4fOGBSoFtk8USmDSno1/FnYRKCH4gcyV3x/vHflFUHlSbh5Aw43YT1wxASJt7lvPq/uSTVw8wVe0bavIW4Gzyk7Fds5fjEi0eyZRtCfAbPlsCQ6aybuZQD870vdT8bxc1sRTdjDRbtFy8PGQqRwR91MhqITLOT7FpptcwltMV9jsGAEXBS6EX754sT3hYLB9OK6INME0rAbHUmq9tvCknIAiH3LIwuJJHHBLEFVeYveTk/00iBHjvn4Yb9MYEPaiTgMcQwRz8khf1aWj/Vz16c2aOjdiZOBEDpxH5h5tJLAMmc8bHfVAhWqdC6hDaulAgHEuo5gcKbkIbWEX/jvuOJq0fO9EYz6eDwwPy6hGGB5pzasJIKUDvXRCh40f5Iy8DOcLJRh+KCxjD9zDqhQQZ8WC+9l1/brckKPw59n0F3Sh3c+nHfyYBBlmsONTmjUZTgRwqg92z+2YQWyUf7g5jmNtBEjLyXtNvn9HkZGXd9YVp7Ps10GklGQiKYZmWUl73KGtsfsRF+/SQ4kRumd4YnlC7b04w6tFRk3HMog+38OVZDwMj40unK4XJWYABJX0t2XySGlXrL8ZNW8xR/iCVsW6+4glxFvTeH5ujPUjQKFb/0bvbZeExeYnRECdDz6u3z/gQYZiUMA8NUNKJuQTzWV9nXyubOWEG9NuJZ4X7oaGE/DtWO3j57r8bcE9KdtV+DpGvKyS+lBrdxL5vlOJ3rX+PqWeIOkxKc85lKT/us8H054bVubQebl0Tc+rvyqZVM7MToHgDj9VwlEbfV6o02Em/J5JUh0GMCoJw6CX5rfHgAIPlhlh/wXRVj8FQcUiTSzDb8lpwXxGO9rNWNfgE9ZRduiXT5LnUYhf6BC5eomyvZ6DcqDABJyWIV7kejbT0TlspicebTzP/kMyTPrGM9TchZjMdv6bPp8jRf4lUGWBjD4i1Diris5aM="
diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml
index 59f046cc9f..97b6f7f8f6 100644
--- a/.github/workflows/CI.yml
+++ b/.github/workflows/CI.yml
@@ -25,7 +25,6 @@ jobs:
           - AdaptiveLoss
           - Logging
           - Forward
-          - NeuralAdapter
           - DGM
         version:
           - "1"
diff --git a/docs/Project.toml b/docs/Project.toml
index 2c81ad0d60..481298e001 100644
--- a/docs/Project.toml
+++ b/docs/Project.toml
@@ -1,5 +1,6 @@
 [deps]
 AdvancedHMC = "0bf59076-c3b1-5ca4-86bd-e02cd72cde3d"
+ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66"
 Cubature = "667455a9-e2ce-5579-9412-b964f529a492"
 DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
@@ -7,8 +8,11 @@ Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 DomainSets = "5b8099bc-c8ec-5219-889f-1d9e522a28bf"
 Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
 Integrals = "de52edbc-65ea-441a-8357-d3a637375a31"
+LineSearches = "d3d80556-e9d4-5f37-9878-2ab0fcc64255"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Lux = "b2108857-7c20-44ae-9111-449ecde12c47"
+LuxCUDA = "d0bbae9a-e099-4d5b-a835-1c6931763bda"
+MethodOfLines = "94925ecb-adb7-4558-8ed8-f975c56a0bf4"
 ModelingToolkit = "961ee093-0014-501f-94e3-6117800e7a78"
 MonteCarloMeasurements = "0987c9cc-fe09-11e8-30f0-b96dd679fdca"
 NeuralPDE = "315f7962-48a3-4962-8226-d0f33b1235f0"
@@ -19,11 +23,13 @@ OptimizationPolyalgorithms = "500b13db-7e66-49ce-bda4-eed966be6282"
 OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
 Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
 QuasiMonteCarlo = "8a4e6c94-4038-4cdc-81c3-7e6ffdb2a71b"
+Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Roots = "f2b01f46-fcfa-551c-844a-d8ac1e96c665"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 
 [compat]
 AdvancedHMC = "0.6"
+ComponentArrays = "0.15"
 Cubature = "1.5"
 DiffEqBase = "6.106"
 Distributions = "0.25"
@@ -31,7 +37,10 @@ Documenter = "1"
 DomainSets = "0.6, 0.7"
 Flux = "0.14"
 Integrals = "4"
+LineSearches = "7.2"
 Lux = "0.5"
+LuxCUDA = "0.3"
+MethodOfLines = "0.10"
 ModelingToolkit = "8.33"
 MonteCarloMeasurements = "1"
 NeuralPDE = "5.3"
@@ -42,5 +51,6 @@ OptimizationPolyalgorithms = "0.2"
 OrdinaryDiffEq = "6.31"
 Plots = "1.36"
 QuasiMonteCarlo = "0.3"
+Random = "1"
 Roots = "2.0"
 SpecialFunctions = "2.1"
diff --git a/docs/src/examples/3rd.md b/docs/src/examples/3rd.md
index 59c68644e2..e64358e177 100644
--- a/docs/src/examples/3rd.md
+++ b/docs/src/examples/3rd.md
@@ -47,7 +47,7 @@ callback = function (p, l)
     return false
 end
 
-res = Optimization.solve(prob, Adam(0.01); callback = callback, maxiters = 2000)
+res = Optimization.solve(prob, OptimizationOptimisers.Adam(0.01); maxiters = 2000)
 phi = discretization.phi
 ```
 
@@ -67,5 +67,3 @@ x_plot = collect(xs)
 plot(x_plot, u_real, title = "real")
 plot!(x_plot, u_predict, title = "predict")
 ```
-
-![hodeplot](https://user-images.githubusercontent.com/12683885/90276340-69bc3e00-de6c-11ea-89a7-7d291123a38b.png)
diff --git a/docs/src/examples/heterogeneous.md b/docs/src/examples/heterogeneous.md
index ad4b72b2ec..286e085be3 100644
--- a/docs/src/examples/heterogeneous.md
+++ b/docs/src/examples/heterogeneous.md
@@ -45,5 +45,5 @@ callback = function (p, l)
     return false
 end
 
-res = Optimization.solve(prob, BFGS(); callback = callback, maxiters = 100)
+res = Optimization.solve(prob, BFGS(); maxiters = 100)
 ```
diff --git a/docs/src/examples/ks.md b/docs/src/examples/ks.md
index bb2b1064f5..55f75f825d 100644
--- a/docs/src/examples/ks.md
+++ b/docs/src/examples/ks.md
@@ -72,7 +72,7 @@ callback = function (p, l)
 end
 
 opt = OptimizationOptimJL.BFGS()
-res = Optimization.solve(prob, opt; callback = callback, maxiters = 2000)
+res = Optimization.solve(prob, opt; maxiters = 2000)
 phi = discretization.phi
 ```
 
@@ -93,5 +93,3 @@ p2 = plot(xs, u_real, title = "analytic")
 p3 = plot(xs, diff_u, title = "error")
 plot(p1, p2, p3)
 ```
-
-![plotks](https://user-images.githubusercontent.com/12683885/91025889-a6253200-e602-11ea-8f61-8e6e2488e025.png)
diff --git a/docs/src/examples/linear_parabolic.md b/docs/src/examples/linear_parabolic.md
index 0ae1432763..60494a5c9a 100644
--- a/docs/src/examples/linear_parabolic.md
+++ b/docs/src/examples/linear_parabolic.md
@@ -23,10 +23,10 @@ w(t, 1) = \frac{e^{\lambda_1} cos(\frac{x}{a})-e^{\lambda_2}cos(\frac{x}{a})}{\l
 
 with a physics-informed neural network.
 
-```@example
-using NeuralPDE, Lux, ModelingToolkit, Optimization, OptimizationOptimJL
+```@example linear_parabolic
+using NeuralPDE, Lux, ModelingToolkit, Optimization, OptimizationOptimisers, OptimizationOptimJL, LineSearches
 using Plots
-import ModelingToolkit: Interval, infimum, supremum
+using ModelingToolkit: Interval, infimum, supremum
 
 @parameters t, x
 @variables u(..), w(..)
@@ -71,7 +71,7 @@ input_ = length(domains)
 n = 15
 chain = [Lux.Chain(Dense(input_, n, Lux.σ), Dense(n, n, Lux.σ), Dense(n, 1)) for _ in 1:2]
 
-strategy = QuadratureTraining()
+strategy = StochasticTraining(500)
 discretization = PhysicsInformedNN(chain, strategy)
 
 @named pdesystem = PDESystem(eqs, bcs, domains, [t, x], [u(t, x), w(t, x)])
@@ -85,14 +85,14 @@ global iteration = 0
 callback = function (p, l)
     if iteration % 10 == 0
         println("loss: ", l)
-        println("pde_losses: ", map(l_ -> l_(p), pde_inner_loss_functions))
-        println("bcs_losses: ", map(l_ -> l_(p), bcs_inner_loss_functions))
+        println("pde_losses: ", map(l_ -> l_(p.u), pde_inner_loss_functions))
+        println("bcs_losses: ", map(l_ -> l_(p.u), bcs_inner_loss_functions))
     end
     global iteration += 1
     return false
 end
 
-res = Optimization.solve(prob, BFGS(); callback = callback, maxiters = 5000)
+res = Optimization.solve(prob, OptimizationOptimisers.Adam(1e-2); maxiters = 10000)
 
 phi = discretization.phi
 
@@ -105,14 +105,19 @@ analytic_sol_func(t, x) = [u_analytic(t, x), w_analytic(t, x)]
 u_real = [[analytic_sol_func(t, x)[i] for t in ts for x in xs] for i in 1:2]
 u_predict = [[phi[i]([t, x], minimizers_[i])[1] for t in ts for x in xs] for i in 1:2]
 diff_u = [abs.(u_real[i] .- u_predict[i]) for i in 1:2]
+ps = []
 for i in 1:2
     p1 = plot(ts, xs, u_real[i], linetype = :contourf, title = "u$i, analytic")
     p2 = plot(ts, xs, u_predict[i], linetype = :contourf, title = "predict")
     p3 = plot(ts, xs, diff_u[i], linetype = :contourf, title = "error")
-    plot(p1, p2, p3)
-    savefig("sol_u$i")
+    push!(ps, plot(p1, p2, p3))
 end
 ```
 
-![linear_parabolic_sol_u1](https://user-images.githubusercontent.com/26853713/125745625-49c73760-0522-4ed4-9bdd-bcc567c9ace3.png)
-![linear_parabolic_sol_u2](https://user-images.githubusercontent.com/26853713/125745637-b12e1d06-e27b-46fe-89f3-076d415fcd7e.png)
+```@example linear_parabolic
+ps[1]
+```
+
+```@example linear_parabolic
+ps[2]
+```
diff --git a/docs/src/examples/nonlinear_elliptic.md b/docs/src/examples/nonlinear_elliptic.md
index db8d6e228f..155330b2bc 100644
--- a/docs/src/examples/nonlinear_elliptic.md
+++ b/docs/src/examples/nonlinear_elliptic.md
@@ -26,10 +26,10 @@ where k is a root of the algebraic (transcendental) equation f(k) = g(k).
 
 This is done using a derivative neural network approximation.
 
-```@example
+```@example nonlinear_elliptic
 using NeuralPDE, Lux, ModelingToolkit, Optimization, OptimizationOptimJL, Roots
 using Plots
-import ModelingToolkit: Interval, infimum, supremum
+using ModelingToolkit: Interval, infimum, supremum
 
 @parameters x, y
 Dx = Differential(x)
@@ -79,7 +79,7 @@ input_ = length(domains)
 n = 15
 chain = [Lux.Chain(Dense(input_, n, Lux.σ), Dense(n, n, Lux.σ), Dense(n, 1)) for _ in 1:6] # 1:number of @variables
 
-strategy = QuadratureTraining()
+strategy = GridTraining(0.01)
 discretization = PhysicsInformedNN(chain, strategy)
 
 vars = [u(x, y), w(x, y), Dxu(x, y), Dyu(x, y), Dxw(x, y), Dyw(x, y)]
@@ -87,10 +87,6 @@ vars = [u(x, y), w(x, y), Dxu(x, y), Dyu(x, y), Dxw(x, y), Dyw(x, y)]
 prob = NeuralPDE.discretize(pdesystem, discretization)
 sym_prob = NeuralPDE.symbolic_discretize(pdesystem, discretization)
 
-strategy = NeuralPDE.QuadratureTraining()
-discretization = PhysicsInformedNN(chain, strategy)
-sym_prob = NeuralPDE.symbolic_discretize(pdesystem, discretization)
-
 pde_inner_loss_functions = sym_prob.loss_functions.pde_loss_functions
 bcs_inner_loss_functions = sym_prob.loss_functions.bc_loss_functions[1:6]
 aprox_derivative_loss_functions = sym_prob.loss_functions.bc_loss_functions[7:end]
@@ -99,15 +95,15 @@ global iteration = 0
 callback = function (p, l)
     if iteration % 10 == 0
         println("loss: ", l)
-        println("pde_losses: ", map(l_ -> l_(p), pde_inner_loss_functions))
-        println("bcs_losses: ", map(l_ -> l_(p), bcs_inner_loss_functions))
-        println("der_losses: ", map(l_ -> l_(p), aprox_derivative_loss_functions))
+        println("pde_losses: ", map(l_ -> l_(p.u), pde_inner_loss_functions))
+        println("bcs_losses: ", map(l_ -> l_(p.u), bcs_inner_loss_functions))
+        println("der_losses: ", map(l_ -> l_(p.u), aprox_derivative_loss_functions))
     end
     global iteration += 1
     return false
 end
 
-res = Optimization.solve(prob, BFGS(); callback = callback, maxiters = 5000)
+res = Optimization.solve(prob, BFGS(); maxiters = 100)
 
 phi = discretization.phi
 
@@ -120,14 +116,19 @@ analytic_sol_func(x, y) = [u_analytic(x, y), w_analytic(x, y)]
 u_real = [[analytic_sol_func(x, y)[i] for x in xs for y in ys] for i in 1:2]
 u_predict = [[phi[i]([x, y], minimizers_[i])[1] for x in xs for y in ys] for i in 1:2]
 diff_u = [abs.(u_real[i] .- u_predict[i]) for i in 1:2]
+ps = []
 for i in 1:2
     p1 = plot(xs, ys, u_real[i], linetype = :contourf, title = "u$i, analytic")
     p2 = plot(xs, ys, u_predict[i], linetype = :contourf, title = "predict")
     p3 = plot(xs, ys, diff_u[i], linetype = :contourf, title = "error")
-    plot(p1, p2, p3)
-    savefig("non_linear_elliptic_sol_u$i")
+    push!(ps, plot(p1, p2, p3))
 end
 ```
 
-![non_linear_elliptic_sol_u1](https://user-images.githubusercontent.com/26853713/125745550-0b667c10-b09a-4659-a543-4f7a7e025d6c.png)
-![non_linear_elliptic_sol_u2](https://user-images.githubusercontent.com/26853713/125745571-45a04739-7838-40ce-b979-43b88d149028.png)
+```@example nonlinear_elliptic
+ps[1]
+```
+
+```@example nonlinear_elliptic
+ps[2]
+```
diff --git a/docs/src/examples/nonlinear_hyperbolic.md b/docs/src/examples/nonlinear_hyperbolic.md
index 3f6896599d..60203b1610 100644
--- a/docs/src/examples/nonlinear_hyperbolic.md
+++ b/docs/src/examples/nonlinear_hyperbolic.md
@@ -32,11 +32,11 @@ where k is a root of the algebraic (transcendental) equation f(k) = g(k), j0 and
 
 We solve this with Neural:
 
-```@example
-using NeuralPDE, Lux, ModelingToolkit, Optimization, OptimizationOptimJL, Roots
+```@example nonlinear_hyperbolic
+using NeuralPDE, Lux, ModelingToolkit, Optimization, OptimizationOptimJL, Roots, LineSearches
 using SpecialFunctions
 using Plots
-import ModelingToolkit: Interval, infimum, supremum
+using ModelingToolkit: Interval, infimum, supremum
 
 @parameters t, x
 @variables u(..), w(..)
@@ -94,12 +94,12 @@ bcs_inner_loss_functions = sym_prob.loss_functions.bc_loss_functions
 
 callback = function (p, l)
     println("loss: ", l)
-    println("pde_losses: ", map(l_ -> l_(p), pde_inner_loss_functions))
-    println("bcs_losses: ", map(l_ -> l_(p), bcs_inner_loss_functions))
+    println("pde_losses: ", map(l_ -> l_(p.u), pde_inner_loss_functions))
+    println("bcs_losses: ", map(l_ -> l_(p.u), bcs_inner_loss_functions))
     return false
 end
 
-res = Optimization.solve(prob, BFGS(); callback = callback, maxiters = 1000)
+res = Optimization.solve(prob, BFGS(linesearch = BackTracking()); maxiters = 200)
 
 phi = discretization.phi
 
@@ -112,14 +112,20 @@ analytic_sol_func(t, x) = [u_analytic(t, x), w_analytic(t, x)]
 u_real = [[analytic_sol_func(t, x)[i] for t in ts for x in xs] for i in 1:2]
 u_predict = [[phi[i]([t, x], minimizers_[i])[1] for t in ts for x in xs] for i in 1:2]
 diff_u = [abs.(u_real[i] .- u_predict[i]) for i in 1:2]
+ps = []
 for i in 1:2
     p1 = plot(ts, xs, u_real[i], linetype = :contourf, title = "u$i, analytic")
     p2 = plot(ts, xs, u_predict[i], linetype = :contourf, title = "predict")
     p3 = plot(ts, xs, diff_u[i], linetype = :contourf, title = "error")
-    plot(p1, p2, p3)
-    savefig("nonlinear_hyperbolic_sol_u$i")
+    push!(ps, plot(p1, p2, p3))
 end
 ```
 
-![nonlinear_hyperbolic_sol_u1](https://user-images.githubusercontent.com/26853713/126457614-d19e7a4d-f9e3-4e78-b8ae-1e58114a744e.png)
-![nonlinear_hyperbolic_sol_u2](https://user-images.githubusercontent.com/26853713/126457617-ee26c587-a97f-4a2e-b6b7-b326b1f117af.png)
+
+```@example nonlinear_hyperbolic
+ps[1]
+```
+
+```@example nonlinear_hyperbolic
+ps[2]
+```
diff --git a/docs/src/examples/wave.md b/docs/src/examples/wave.md
index 6d99f0fa56..954b15cc9a 100644
--- a/docs/src/examples/wave.md
+++ b/docs/src/examples/wave.md
@@ -17,7 +17,7 @@ Further, the solution of this equation with the given boundary conditions is pre
 
 ```@example wave
 using NeuralPDE, Lux, Optimization, OptimizationOptimJL
-import ModelingToolkit: Interval
+using ModelingToolkit: Interval
 
 @parameters t, x
 @variables u(..)
@@ -81,8 +81,6 @@ p3 = plot(ts, xs, diff_u, linetype = :contourf, title = "error");
 plot(p1, p2, p3)
 ```
 
-![waveplot](https://user-images.githubusercontent.com/12683885/101984293-74a7a380-3c91-11eb-8e78-72a50d88e3f8.png)
-
 ## 1D Damped Wave Equation with Dirichlet boundary conditions
 
 Now let's solve the 1-dimensional wave equation with damping.
@@ -101,7 +99,7 @@ with grid discretization `dx = 0.05` and physics-informed neural networks. Here,
 ```@example wave2
 using NeuralPDE, Lux, ModelingToolkit, Optimization, OptimizationOptimJL
 using Plots, Printf
-import ModelingToolkit: Interval, infimum, supremum
+using ModelingToolkit: Interval, infimum, supremum
 
 @parameters t, x
 @variables u(..) Dxu(..) Dtu(..) O1(..) O2(..)
@@ -159,14 +157,14 @@ bcs_inner_loss_functions = sym_prob.loss_functions.bc_loss_functions
 
 callback = function (p, l)
     println("loss: ", l)
-    println("pde_losses: ", map(l_ -> l_(p), pde_inner_loss_functions))
-    println("bcs_losses: ", map(l_ -> l_(p), bcs_inner_loss_functions))
+    println("pde_losses: ", map(l_ -> l_(p.u), pde_inner_loss_functions))
+    println("bcs_losses: ", map(l_ -> l_(p.u), bcs_inner_loss_functions))
     return false
 end
 
-res = Optimization.solve(prob, BFGS(); callback = callback, maxiters = 2000)
+res = Optimization.solve(prob, BFGS(); maxiters = 2000)
 prob = remake(prob, u0 = res.u)
-res = Optimization.solve(prob, BFGS(); callback = callback, maxiters = 2000)
+res = Optimization.solve(prob, BFGS(); maxiters = 2000)
 
 phi = discretization.phi[1]
 
@@ -214,11 +212,3 @@ p2 = plot(ts, xs, u_predict, linetype = :contourf, title = "predict");
 p3 = plot(ts, xs, diff_u, linetype = :contourf, title = "error");
 plot(p1, p2, p3)
 ```
-
-We can see the results here:
-
-![Damped_wave_sol_adaptive_u](https://user-images.githubusercontent.com/12683885/149665332-d4daf7d0-682e-4933-a2b4-34f403881afb.png)
-
-Plotted as a line, one can see the analytical solution and the prediction here:
-
-![1Dwave_damped_adaptive](https://user-images.githubusercontent.com/12683885/149665327-69d04c01-2240-45ea-981e-a7b9412a3b58.gif)
diff --git a/docs/src/tutorials/constraints.md b/docs/src/tutorials/constraints.md
index 68aa391af7..59f389b172 100644
--- a/docs/src/tutorials/constraints.md
+++ b/docs/src/tutorials/constraints.md
@@ -21,9 +21,9 @@ p(-2.2) = p(2.2) = 0
 with Physics-Informed Neural Networks.
 
 ```@example fokkerplank
-using NeuralPDE, Lux, ModelingToolkit, Optimization, OptimizationOptimJL
+using NeuralPDE, Lux, ModelingToolkit, Optimization, OptimizationOptimJL, LineSearches
 using Integrals, Cubature
-import ModelingToolkit: Interval, infimum, supremum
+using ModelingToolkit: Interval, infimum, supremum
 # the example is taken from this article https://arxiv.org/abs/1910.10503
 @parameters x
 @variables p(..)
@@ -78,15 +78,13 @@ aprox_derivative_loss_functions = sym_prob.loss_functions.bc_loss_functions
 
 cb_ = function (p, l)
     println("loss: ", l)
-    println("pde_losses: ", map(l_ -> l_(p), pde_inner_loss_functions))
-    println("bcs_losses: ", map(l_ -> l_(p), bcs_inner_loss_functions))
-    println("additional_loss: ", norm_loss_function(phi, p, nothing))
+    println("pde_losses: ", map(l_ -> l_(p.u), pde_inner_loss_functions))
+    println("bcs_losses: ", map(l_ -> l_(p.u), bcs_inner_loss_functions))
+    println("additional_loss: ", norm_loss_function(phi, p.u, nothing))
     return false
 end
 
-res = Optimization.solve(prob, LBFGS(), callback = cb_, maxiters = 400)
-prob = remake(prob, u0 = res.u)
-res = Optimization.solve(prob, BFGS(), callback = cb_, maxiters = 2000)
+res = Optimization.solve(prob, BFGS(linesearch = BackTracking()), callback = cb_, maxiters = 600)
 ```
 
 And some analysis:
@@ -103,5 +101,3 @@ u_predict = [first(phi(x, res.u)) for x in xs]
 plot(xs, u_real, label = "analytic")
 plot!(xs, u_predict, label = "predict")
 ```
-
-![fp](https://user-images.githubusercontent.com/12683885/129405830-3d00c24e-adf1-443b-aa36-6af0e5305821.png)
diff --git a/docs/src/tutorials/dae.md b/docs/src/tutorials/dae.md
index 24a7880a6e..de83542c63 100644
--- a/docs/src/tutorials/dae.md
+++ b/docs/src/tutorials/dae.md
@@ -15,14 +15,14 @@ Let's solve a simple DAE system:
 
 ```@example dae
 using NeuralPDE 
-using Random, Flux
-using OrdinaryDiffEq, Optimisers, Statistics
-import Lux, OptimizationOptimisers, OptimizationOptimJL
+using Random
+using OrdinaryDiffEq, Statistics
+using Lux, OptimizationOptimisers
 
 example = (du, u, p, t) -> [cos(2pi * t) - du[1], u[2] + cos(2pi * t) - du[2]]
 u₀ = [1.0, -1.0]
 du₀ = [0.0, 0.0]
-tspan = (0.0f0, 1.0f0)
+tspan = (0.0, 1.0)
 ```
 
 Differential_vars is a logical array which declares which variables are the differential (non-algebraic) vars
@@ -33,13 +33,10 @@ differential_vars = [true, false]
 
 ```@example dae
 prob = DAEProblem(example, du₀, u₀, tspan; differential_vars = differential_vars)
-chain = Flux.Chain(Dense(1, 15, cos), Dense(15, 15, sin), Dense(15, 2))
+chain = Lux.Chain(Lux.Dense(1, 15, cos), Lux.Dense(15, 15, sin), Lux.Dense(15, 2))
 opt = OptimizationOptimisers.Adam(0.1)
 alg = NNDAE(chain, opt; autodiff = false)
-
-sol = solve(prob,
-    alg, verbose = false, dt = 1 / 100.0f0,
-    maxiters = 3000, abstol = 1.0f-10)
+sol = solve(prob, alg, verbose = true, dt = 1 / 100.0, maxiters = 3000, abstol = 1e-10)
 ```
 
 Now lets compare the predictions from the learned network with the ground truth which we can obtain by numerically solving the DAE.
@@ -50,8 +47,7 @@ function example1(du, u, p, t)
     du[2] = u[2] + cos(2pi * t)
     nothing
 end
-M = [1.0 0
-    0 0]
+M = [1.0 0.0; 0.0 0.0]
 f = ODEFunction(example1, mass_matrix = M)
 prob_mm = ODEProblem(f, u₀, tspan)
 ground_sol = solve(prob_mm, Rodas5(), reltol = 1e-8, abstol = 1e-8)
diff --git a/docs/src/tutorials/derivative_neural_network.md b/docs/src/tutorials/derivative_neural_network.md
index 3432d58688..82ad99b986 100644
--- a/docs/src/tutorials/derivative_neural_network.md
+++ b/docs/src/tutorials/derivative_neural_network.md
@@ -53,9 +53,9 @@ using the second numeric derivative `Dt(Dtu1(t,x))`.
 
 ```@example derivativenn
 using NeuralPDE, Lux, ModelingToolkit
-using Optimization, OptimizationOptimisers, OptimizationOptimJL
+using Optimization, OptimizationOptimisers, OptimizationOptimJL, LineSearches
 using Plots
-import ModelingToolkit: Interval, infimum, supremum
+using ModelingToolkit: Interval, infimum, supremum
 
 @parameters t, x
 Dt = Differential(t)
@@ -71,8 +71,6 @@ bcs_ = [u1(0.0, x) ~ sin(pi * x),
     u2(0.0, x) ~ cos(pi * x),
     Dt(u1(0, x)) ~ -sin(pi * x),
     Dt(u2(0, x)) ~ -cos(pi * x),
-    #Dtu1(0,x) ~ -sin(pi*x),
-    # Dtu2(0,x) ~ -cos(pi*x),
     u1(t, 0.0) ~ 0.0,
     u2(t, 0.0) ~ exp(-t),
     u1(t, 1.0) ~ 0.0,
@@ -93,9 +91,8 @@ input_ = length(domains)
 n = 15
 chain = [Lux.Chain(Dense(input_, n, Lux.σ), Dense(n, n, Lux.σ), Dense(n, 1)) for _ in 1:7]
 
-training_strategy = NeuralPDE.QuadratureTraining()
-discretization = NeuralPDE.PhysicsInformedNN(chain,
-                                             training_strategy)
+training_strategy = NeuralPDE.QuadratureTraining(; batch = 200, reltol = 1e-6, abstol = 1e-6)
+discretization = NeuralPDE.PhysicsInformedNN(chain, training_strategy)
 
 vars = [u1(t, x), u2(t, x), u3(t, x), Dxu1(t, x), Dtu1(t, x), Dxu2(t, x), Dtu2(t, x)]
 @named pdesystem = PDESystem(eqs_, bcs__, domains, [t, x], vars)
@@ -108,15 +105,15 @@ aprox_derivative_loss_functions = sym_prob.loss_functions.bc_loss_functions[9:en
 
 callback = function (p, l)
     println("loss: ", l)
-    println("pde_losses: ", map(l_ -> l_(p), pde_inner_loss_functions))
-    println("bcs_losses: ", map(l_ -> l_(p), bcs_inner_loss_functions))
-    println("der_losses: ", map(l_ -> l_(p), aprox_derivative_loss_functions))
+    println("pde_losses: ", map(l_ -> l_(p.u), pde_inner_loss_functions))
+    println("bcs_losses: ", map(l_ -> l_(p.u), bcs_inner_loss_functions))
+    println("der_losses: ", map(l_ -> l_(p.u), aprox_derivative_loss_functions))
     return false
 end
 
-res = Optimization.solve(prob, Adam(0.01); callback = callback, maxiters = 2000)
+res = Optimization.solve(prob, OptimizationOptimisers.Adam(0.01); maxiters = 2000)
 prob = remake(prob, u0 = res.u)
-res = Optimization.solve(prob, BFGS(); callback = callback, maxiters = 10000)
+res = Optimization.solve(prob, LBFGS(linesearch = BackTracking()); maxiters = 200)
 
 phi = discretization.phi
 ```
@@ -136,6 +133,7 @@ Dxu1_real(t, x) = pi * exp(-t) * cos(pi * x)
 Dtu1_real(t, x) = -exp(-t) * sin(pi * x)
 Dxu2_real(t, x) = -pi * exp(-t) * sin(pi * x)
 Dtu2_real(t, x) = -exp(-t) * cos(pi * x)
+
 function analytic_sol_func_all(t, x)
     [u1_real(t, x), u2_real(t, x), u3_real(t, x),
         Dxu1_real(t, x), Dtu1_real(t, x), Dxu2_real(t, x), Dtu2_real(t, x)]
@@ -144,25 +142,40 @@ end
 u_real = [[analytic_sol_func_all(t, x)[i] for t in ts for x in xs] for i in 1:7]
 u_predict = [[phi[i]([t, x], minimizers_[i])[1] for t in ts for x in xs] for i in 1:7]
 diff_u = [abs.(u_real[i] .- u_predict[i]) for i in 1:7]
-
+ps = []
 titles = ["u1", "u2", "u3", "Dtu1", "Dtu2", "Dxu1", "Dxu2"]
 for i in 1:7
     p1 = plot(ts, xs, u_real[i], linetype = :contourf, title = "$(titles[i]), analytic")
     p2 = plot(ts, xs, u_predict[i], linetype = :contourf, title = "predict")
     p3 = plot(ts, xs, diff_u[i], linetype = :contourf, title = "error")
-    plot(p1, p2, p3)
-    savefig("3sol_ub$i")
+    push!(ps, plot(p1, p2, p3))
 end
 ```
 
-![aprNN_sol_u1](https://user-images.githubusercontent.com/12683885/122998551-de79d600-d3b5-11eb-8f5d-59d00178c2ab.png)
+```@example derivativenn
+ps[1]
+```
 
-![aprNN_sol_u2](https://user-images.githubusercontent.com/12683885/122998567-e3d72080-d3b5-11eb-9024-4072f4b66cda.png)
+```@example derivativenn
+ps[2]
+```
 
-![aprNN_sol_u3](https://user-images.githubusercontent.com/12683885/122998578-e6d21100-d3b5-11eb-96a5-f64e5593b35e.png)
+```@example derivativenn
+ps[3]
+```
 
-## Comparison of the second numerical derivative and numerical + neural network derivative
+```@example derivativenn
+ps[4]
+```
+
+```@example derivativenn
+ps[5]
+```
 
-![DDu1](https://user-images.githubusercontent.com/12683885/123113394-3280cb00-d447-11eb-88e3-a8541bbf089f.png)
+```@example derivativenn
+ps[6]
+```
 
-![DDu2](https://user-images.githubusercontent.com/12683885/123113413-36ace880-d447-11eb-8f6a-4c3caa86e359.png)
+```@example derivativenn
+ps[7]
+```
diff --git a/docs/src/tutorials/dgm.md b/docs/src/tutorials/dgm.md
index b20eb30bb9..1917df8f43 100644
--- a/docs/src/tutorials/dgm.md
+++ b/docs/src/tutorials/dgm.md
@@ -7,6 +7,7 @@ Deep Galerkin Method is a meshless deep learning algorithm to solve high dimensi
 In the following example, we demonstrate computing the loss function using Quasi-Random Sampling, a sampling technique that uses quasi-Monte Carlo sampling to generate low discrepancy random sequences in high dimensional spaces.
 
 ### Algorithm
+
 The authors of DGM suggest a network composed of LSTM-type layers that works well for most of the parabolic and quasi-parabolic PDEs.
 
 ```math
@@ -47,13 +48,15 @@ u(t, 1) & = 0
 ```
 
 ### Copy- Pasteable code
+
 ```@example dgm
 using NeuralPDE
 using ModelingToolkit, Optimization, OptimizationOptimisers
-import Lux: tanh, identity
+using Lux: tanh, identity
 using Distributions
-import ModelingToolkit: Interval, infimum, supremum
+using ModelingToolkit: Interval, infimum, supremum
 using MethodOfLines, OrdinaryDiffEq
+using Plots
 
 @parameters x t
 @variables u(..)
@@ -100,16 +103,17 @@ callback = function (p, l)
     return false
 end
 
-res = Optimization.solve(prob, Adam(0.1); callback = callback, maxiters = 100)
+res = Optimization.solve(prob, Adam(0.1); maxiters = 100)
 prob = remake(prob, u0 = res.u)
-res = Optimization.solve(prob, Adam(0.01); callback = callback, maxiters = 500)
+res = Optimization.solve(prob, Adam(0.01); maxiters = 500)
 phi = discretization.phi
 
 u_predict= [first(phi([t, x], res.minimizer)) for t in ts, x in xs]
 
 diff_u = abs.(u_predict .- u_MOL)
+tgrid = 0.0:0.01:1.0
+xgrid = -1.0:0.01:1.0
 
-using Plots
 p1 = plot(tgrid, xgrid, u_MOL', linetype = :contourf, title = "FD");
 p2 = plot(tgrid, xgrid, u_predict', linetype = :contourf, title = "predict");
 p3 = plot(tgrid, xgrid, diff_u', linetype = :contourf, title = "error");
diff --git a/docs/src/tutorials/gpu.md b/docs/src/tutorials/gpu.md
index 5a95043450..feea05c9e1 100644
--- a/docs/src/tutorials/gpu.md
+++ b/docs/src/tutorials/gpu.md
@@ -29,24 +29,28 @@ we must ensure that our initial parameters for the neural network are on the GPU
 is done, then the internal computations will all take place on the GPU. This is done by
 using the `gpu` function on the initial parameters, like:
 
-```julia
-using Lux, ComponentArrays
+```@example gpu
+using Lux, LuxCUDA, ComponentArrays, Random
+const gpud = gpu_device()
+inner = 25
 chain = Chain(Dense(3, inner, Lux.σ),
               Dense(inner, inner, Lux.σ),
               Dense(inner, inner, Lux.σ),
               Dense(inner, inner, Lux.σ),
               Dense(inner, 1))
 ps = Lux.setup(Random.default_rng(), chain)[1]
-ps = ps |> ComponentArray |> gpu .|> Float64
+ps = ps |> ComponentArray |> gpud .|> Float64
 ```
 
 In total, this looks like:
 
-```julia
-using NeuralPDE, Lux, CUDA, Random, ComponentArrays
+```@example gpu
+using NeuralPDE, Lux, LuxCUDA, Random, ComponentArrays
 using Optimization
 using OptimizationOptimisers
 import ModelingToolkit: Interval
+using Plots
+using Printf
 
 @parameters t x y
 @variables u(..)
@@ -84,9 +88,9 @@ chain = Chain(Dense(3, inner, Lux.σ),
               Dense(inner, inner, Lux.σ),
               Dense(inner, 1))
 
-strategy = QuadratureTraining()
+strategy = QuasiRandomTraining(100)
 ps = Lux.setup(Random.default_rng(), chain)[1]
-ps = ps |> ComponentArray |> gpu .|> Float64
+ps = ps |> ComponentArray |> gpud .|> Float64
 discretization = PhysicsInformedNN(chain,
                                    strategy,
                                    init_params = ps)
@@ -100,27 +104,23 @@ callback = function (p, l)
     return false
 end
 
-res = Optimization.solve(prob, Adam(0.01); callback = callback, maxiters = 2500)
+res = Optimization.solve(prob, OptimizationOptimisers.Adam(1e-2); maxiters = 2500)
 ```
 
-We then use the `remake` function to rebuild the PDE problem to start a new
-optimization at the optimized parameters, and continue with a lower learning rate:
+We then use the `remake` function to rebuild the PDE problem to start a new optimization at the optimized parameters, and continue with a lower learning rate:
 
-```julia
+```@example gpu
 prob = remake(prob, u0 = res.u)
-res = Optimization.solve(prob, Adam(0.001); callback = callback, maxiters = 2500)
+res = Optimization.solve(prob, OptimizationOptimisers.Adam(1e-3); callback = callback, maxiters = 2500)
 ```
 
 Finally, we inspect the solution:
 
-```julia
+```@example gpu
 phi = discretization.phi
 ts, xs, ys = [infimum(d.domain):0.1:supremum(d.domain) for d in domains]
 u_real = [analytic_sol_func(t, x, y) for t in ts for x in xs for y in ys]
-u_predict = [first(Array(phi(gpu([t, x, y]), res.u))) for t in ts for x in xs for y in ys]
-
-using Plots
-using Printf
+u_predict = [first(Array(phi([t, x, y], res.u))) for t in ts for x in xs for y in ys]
 
 function plot_(res)
     # Animate
@@ -128,7 +128,7 @@ function plot_(res)
         @info "Animating frame $i..."
         u_real = reshape([analytic_sol_func(t, x, y) for x in xs for y in ys],
                          (length(xs), length(ys)))
-        u_predict = reshape([Array(phi(gpu([t, x, y]), res.u))[1] for x in xs for y in ys],
+        u_predict = reshape([Array(phi([t, x, y], res.u))[1] for x in xs for y in ys],
                             length(xs), length(ys))
         u_error = abs.(u_predict .- u_real)
         title = @sprintf("predict, t = %.3f", t)
@@ -145,8 +145,6 @@ end
 plot_(res)
 ```
 
-![3pde](https://user-images.githubusercontent.com/12683885/129949743-9471d230-c14f-4105-945f-6bc52677d40e.gif)
-
 ## Performance benchmarks
 
 Here are some performance benchmarks for 2d-pde with various number of input points and the
@@ -155,7 +153,6 @@ runtime with GPU and CPU.
 
 ```julia
 julia> CUDA.device()
-
 ```
 
 ![image](https://user-images.githubusercontent.com/12683885/110297207-49202500-8004-11eb-9e45-d4cb28045d87.png)
diff --git a/docs/src/tutorials/integro_diff.md b/docs/src/tutorials/integro_diff.md
index 3eb2323ba0..e977b87af4 100644
--- a/docs/src/tutorials/integro_diff.md
+++ b/docs/src/tutorials/integro_diff.md
@@ -45,8 +45,9 @@ u(0) = 0
 ```
 
 ```@example integro
-using NeuralPDE, Flux, ModelingToolkit, Optimization, OptimizationOptimJL, DomainSets
-import ModelingToolkit: Interval, infimum, supremum
+using NeuralPDE, Lux, ModelingToolkit, Optimization, OptimizationOptimJL, DomainSets
+using ModelingToolkit: Interval, infimum, supremum
+using Plots
 
 @parameters t
 @variables i(..)
@@ -55,7 +56,7 @@ Ii = Integral(t in DomainSets.ClosedInterval(0, t))
 eq = Di(i(t)) + 2 * i(t) + 5 * Ii(i(t)) ~ 1
 bcs = [i(0.0) ~ 0.0]
 domains = [t ∈ Interval(0.0, 2.0)]
-chain = Chain(Dense(1, 15, Flux.σ), Dense(15, 1)) |> f64
+chain = Lux.Chain(Lux.Dense(1, 15, Lux.σ), Lux.Dense(15, 1))
 
 strategy_ = QuadratureTraining()
 discretization = PhysicsInformedNN(chain,
@@ -66,7 +67,7 @@ callback = function (p, l)
     println("Current loss is: $l")
     return false
 end
-res = Optimization.solve(prob, BFGS(); callback = callback, maxiters = 100)
+res = Optimization.solve(prob, BFGS(); maxiters = 100)
 ```
 
 Plotting the final solution and analytical solution
@@ -78,9 +79,6 @@ u_predict = [first(phi([t], res.u)) for t in ts]
 
 analytic_sol_func(t) = 1 / 2 * (exp(-t)) * (sin(2 * t))
 u_real = [analytic_sol_func(t) for t in ts]
-using Plots
 plot(ts, u_real, label = "Analytical Solution")
 plot!(ts, u_predict, label = "PINN Solution")
 ```
-
-![IDE](https://user-images.githubusercontent.com/12683885/129749371-18b44bbc-18c8-49c5-bf30-0cd97ecdd977.png)
diff --git a/docs/src/tutorials/low_level.md b/docs/src/tutorials/low_level.md
index aad1347971..33b605e77d 100644
--- a/docs/src/tutorials/low_level.md
+++ b/docs/src/tutorials/low_level.md
@@ -13,8 +13,8 @@ u(t, -1) = u(t, 1) = 0 \, ,
 with Physics-Informed Neural Networks. Here is an example of using the low-level API:
 
 ```@example low_level
-using NeuralPDE, Lux, ModelingToolkit, Optimization, OptimizationOptimJL
-import ModelingToolkit: Interval, infimum, supremum
+using NeuralPDE, Lux, ModelingToolkit, Optimization, OptimizationOptimJL, LineSearches
+using ModelingToolkit: Interval, infimum, supremum
 
 @parameters t, x
 @variables u(..)
@@ -37,7 +37,7 @@ domains = [t ∈ Interval(0.0, 1.0),
 
 # Neural network
 chain = Lux.Chain(Dense(2, 16, Lux.σ), Dense(16, 16, Lux.σ), Dense(16, 1))
-strategy = NeuralPDE.QuadratureTraining()
+strategy = NeuralPDE.QuadratureTraining(; abstol = 1e-6, reltol = 1e-6, batch = 200)
 
 indvars = [t, x]
 depvars = [u(t, x)]
@@ -53,8 +53,8 @@ bc_loss_functions = sym_prob.loss_functions.bc_loss_functions
 
 callback = function (p, l)
     println("loss: ", l)
-    println("pde_losses: ", map(l_ -> l_(p), pde_loss_functions))
-    println("bcs_losses: ", map(l_ -> l_(p), bc_loss_functions))
+    println("pde_losses: ", map(l_ -> l_(p.u), pde_loss_functions))
+    println("bcs_losses: ", map(l_ -> l_(p.u), bc_loss_functions))
     return false
 end
 
@@ -67,8 +67,7 @@ end
 f_ = OptimizationFunction(loss_function, Optimization.AutoZygote())
 prob = Optimization.OptimizationProblem(f_, sym_prob.flat_init_params)
 
-res = Optimization.solve(prob, OptimizationOptimJL.BFGS(); callback = callback,
-                         maxiters = 2000)
+res = Optimization.solve(prob, BFGS(linesearch = BackTracking()); maxiters = 3000)
 ```
 
 And some analysis:
@@ -87,7 +86,3 @@ p2 = plot(xs, u_predict[11], title = "t = 0.5");
 p3 = plot(xs, u_predict[end], title = "t = 1");
 plot(p1, p2, p3)
 ```
-
-![burgers](https://user-images.githubusercontent.com/12683885/90984874-a0870800-e580-11ea-9fd4-af8a4e3c523e.png)
-
-![burgers2](https://user-images.githubusercontent.com/12683885/90984856-8c430b00-e580-11ea-9206-1a88ebd24ca0.png)
diff --git a/docs/src/tutorials/neural_adapter.md b/docs/src/tutorials/neural_adapter.md
index 762a5108f6..a56e30a269 100644
--- a/docs/src/tutorials/neural_adapter.md
+++ b/docs/src/tutorials/neural_adapter.md
@@ -1,9 +1,5 @@
 # Transfer Learning with Neural Adapter
 
-!!! warning
-    
-    This documentation page is out of date.
-
 Transfer learning is a machine learning technique where a model trained on one task is re-purposed on a second related task.
 
 `neural_adapter` is the method that trains a neural network using the results from an already obtained prediction.
@@ -17,9 +13,10 @@ Using the example of 2D Poisson equation, it is shown how, using the method neur
 
 ![image](https://user-images.githubusercontent.com/12683885/127149639-c2a8066f-9a25-4889-b313-5d4403567300.png)
 
-```julia
-using NeuralPDE, Lux, ModelingToolkit, Optimization, OptimizationOptimJL, DiffEqBase
-import ModelingToolkit: Interval, infimum, supremum
+```@example neural_adapter
+using NeuralPDE, Lux, ModelingToolkit, Optimization, OptimizationOptimisers
+using ModelingToolkit: Interval, infimum, supremum
+using Random, ComponentArrays
 
 @parameters x y
 @variables u(..)
@@ -33,57 +30,58 @@ eq = Dxx(u(x, y)) + Dyy(u(x, y)) ~ -sin(pi * x) * sin(pi * y)
 bcs = [u(0, y) ~ 0.0, u(1, y) ~ -sin(pi * 1) * sin(pi * y),
     u(x, 0) ~ 0.0, u(x, 1) ~ -sin(pi * x) * sin(pi * 1)]
 # Space and time domains
-domains = [x ∈ Interval(0.0, 1.0),
-    y ∈ Interval(0.0, 1.0)]
-quadrature_strategy = NeuralPDE.QuadratureTraining(reltol = 1e-2, abstol = 1e-2,
+domains = [x ∈ Interval(0.0, 1.0), y ∈ Interval(0.0, 1.0)]
+quadrature_strategy = NeuralPDE.QuadratureTraining(reltol = 1e-3, abstol = 1e-6,
                                                    maxiters = 50, batch = 100)
 inner = 8
 af = Lux.tanh
-chain1 = Chain(Dense(2, inner, af),
-               Dense(inner, inner, af),
-               Dense(inner, 1))
+chain1 = Lux.Chain(Lux.Dense(2, inner, af),
+                   Lux.Dense(inner, inner, af),
+                   Lux.Dense(inner, 1))
 
-discretization = NeuralPDE.PhysicsInformedNN(chain1,
-                                             quadrature_strategy)
+discretization = NeuralPDE.PhysicsInformedNN(chain1, quadrature_strategy)
 
 @named pde_system = PDESystem(eq, bcs, domains, [x, y], [u(x, y)])
 prob = NeuralPDE.discretize(pde_system, discretization)
 sym_prob = NeuralPDE.symbolic_discretize(pde_system, discretization)
 
-res = Optimization.solve(prob, BFGS(); maxiters = 2000)
+callback = function (p, l)
+    println("Current loss is: $l")
+    return false
+end
+
+res = Optimization.solve(prob, OptimizationOptimisers.Adam(5e-3); maxiters = 10000)
 phi = discretization.phi
 
-inner_ = 12
+inner_ = 8
 af = Lux.tanh
 chain2 = Lux.Chain(Dense(2, inner_, af),
                    Dense(inner_, inner_, af),
                    Dense(inner_, inner_, af),
                    Dense(inner_, 1))
-
-init_params2 = Float64.(ComponentArray(Lux.setup(Random.default_rng(), chain)[1]))
+initp, st = Lux.setup(Random.default_rng(), chain2)
+init_params2 = Float64.(ComponentArrays.ComponentArray(initp))
 
 # the rule by which the training will take place is described here in loss function
 function loss(cord, θ)
-    chain2(cord, θ) .- phi(cord, res.u)
+    global st, chain2
+    ch2, st = chain2(cord, θ, st)
+    ch2 .- phi(cord, res.u)
 end
 
 strategy = NeuralPDE.QuadratureTraining()
 
 prob_ = NeuralPDE.neural_adapter(loss, init_params2, pde_system, strategy)
-callback = function (p, l)
-    println("Current loss is: $l")
-    return false
-end
-res_ = Optimization.solve(prob_, BFGS(); callback = callback, maxiters = 1000)
+res_ = Optimization.solve(prob_, OptimizationOptimisers.Adam(5e-3); maxiters = 10000)
 
-phi_ = NeuralPDE.get_phi(chain2)
+phi_ = PhysicsInformedNN(chain2, strategy; init_params = res_.u).phi
 
 xs, ys = [infimum(d.domain):0.01:supremum(d.domain) for d in domains]
 analytic_sol_func(x, y) = (sin(pi * x) * sin(pi * y)) / (2pi^2)
 
 u_predict = reshape([first(phi([x, y], res.u)) for x in xs for y in ys],
                     (length(xs), length(ys)))
-u_predict_ = reshape([first(phi_([x, y], res_.minimizer)) for x in xs for y in ys],
+u_predict_ = reshape([first(phi_([x, y], res_.u)) for x in xs for y in ys],
                      (length(xs), length(ys)))
 u_real = reshape([analytic_sol_func(x, y) for x in xs for y in ys],
                  (length(xs), length(ys)))
@@ -91,16 +89,14 @@ diff_u = u_predict .- u_real
 diff_u_ = u_predict_ .- u_real
 
 using Plots
-p1 = plot(xs, ys, u_predict, linetype = :contourf, title = "first predict");
-p2 = plot(xs, ys, u_predict_, linetype = :contourf, title = "second predict");
-p3 = plot(xs, ys, u_real, linetype = :contourf, title = "analytic");
-p4 = plot(xs, ys, diff_u, linetype = :contourf, title = "error 1");
-p5 = plot(xs, ys, diff_u_, linetype = :contourf, title = "error 2");
+p1 = plot(xs, ys, u_predict, linetype = :contourf, title = "first predict")
+p2 = plot(xs, ys, u_predict_, linetype = :contourf, title = "second predict")
+p3 = plot(xs, ys, u_real, linetype = :contourf, title = "analytic")
+p4 = plot(xs, ys, diff_u, linetype = :contourf, title = "error 1")
+p5 = plot(xs, ys, diff_u_, linetype = :contourf, title = "error 2")
 plot(p1, p2, p3, p4, p5)
 ```
 
-![neural_adapter](https://user-images.githubusercontent.com/12683885/127645487-24226cc0-fff6-4bb9-8509-77cf3e120563.png)
-
 ## Domain decomposition
 
 In this example, we first obtain a prediction of 2D Poisson equation on subdomains. We split up full domain into 10 sub problems by x, and create separate neural networks for each sub interval. If x domain ∈ [x_0, x_end] so, it is decomposed on 10 part: sub x domains = {[x_0, x_1], ... [x_i,x_i+1], ..., x_9,x_end]}.
@@ -108,9 +104,9 @@ And then using the method neural_adapter, we retrain the batch of 10 predictions
 
 ![domain_decomposition](https://user-images.githubusercontent.com/12683885/127149752-a4ecea50-2984-45d8-b0d4-d2eadecf58e7.png)
 
-```julia
-using NeuralPDE, Lux, ModelingToolkit, Optimization, OptimizationOptimJL, DiffEqBase
-import ModelingToolkit: Interval, infimum, supremum
+```@example neural_adapter
+using NeuralPDE, Lux, ModelingToolkit, Optimization, OptimizationOptimisers
+using ModelingToolkit: Interval, infimum, supremum
 
 @parameters x y
 @variables u(..)
@@ -127,8 +123,7 @@ x_0 = 0.0
 x_end = 1.0
 x_domain = Interval(x_0, x_end)
 y_domain = Interval(0.0, 1.0)
-domains = [x ∈ x_domain,
-    y ∈ y_domain]
+domains = [x ∈ x_domain, y ∈ y_domain]
 
 count_decomp = 10
 
@@ -137,8 +132,7 @@ af = Lux.tanh
 inner = 10
 chains = [Lux.Chain(Dense(2, inner, af), Dense(inner, inner, af), Dense(inner, 1))
           for _ in 1:count_decomp]
-init_params = map(c -> Float64.(ComponentArray(Lux.setup(Random.default_rng(), c)[1])),
-                  chains)
+init_params = map(c -> Float64.(ComponentArray(Lux.setup(Random.default_rng(), c)[1])), chains)
 
 xs_ = infimum(x_domain):(1 / count_decomp):supremum(x_domain)
 xs_domain = [(xs_[i], xs_[i + 1]) for i in 1:(length(xs_) - 1)]
@@ -172,9 +166,9 @@ pde_system_map = []
 for i in 1:count_decomp
     println("decomposition $i")
     domains_ = domains_map[i]
-    phi_in(cord) = phis[i - 1](cord, reses[i - 1].minimizer)
+    phi_in(cord) = phis[i - 1](cord, reses[i - 1].u)
     phi_bound(x, y) = phi_in(vcat(x, y))
-    @register phi_bound(x, y)
+    @register_symbolic phi_bound(x, y)
     Base.Broadcast.broadcasted(::typeof(phi_bound), x, y) = phi_bound(x, y)
     bcs_ = create_bcs(domains_[1].domain, phi_bound)
     @named pde_system_ = PDESystem(eq, bcs_, domains_, [x, y], [u(x, y)])
@@ -186,7 +180,7 @@ for i in 1:count_decomp
 
     prob = NeuralPDE.discretize(pde_system_, discretization)
     symprob = NeuralPDE.symbolic_discretize(pde_system_, discretization)
-    res_ = Optimization.solve(prob, BFGS(), maxiters = 1000)
+    res_ = Optimization.solve(prob, OptimizationOptimisers.Adam(5e-3); maxiters = 10000)
     phi = discretization.phi
     push!(reses, res_)
     push!(phis, phi)
@@ -207,7 +201,7 @@ function compose_result(dx)
     end
     for x_ in xs
         i = index_of_interval(x_)
-        u_predict_sub = [first(phis[i]([x_, y], reses[i].minimizer)) for y in ys]
+        u_predict_sub = [first(phis[i]([x_, y], reses[i].u)) for y in ys]
         u_real_sub = [analytic_sol_func(x_, y) for y in ys]
         diff_u_sub = abs.(u_predict_sub .- u_real_sub)
         append!(u_predict_array, u_predict_sub)
@@ -218,6 +212,7 @@ function compose_result(dx)
     diff_u = reshape(diff_u_array, (length(xs), length(ys)))
     u_predict, diff_u
 end
+
 dx = 0.01
 u_predict, diff_u = compose_result(dx)
 
@@ -229,12 +224,17 @@ chain2 = Lux.Chain(Dense(2, inner_, af),
                    Dense(inner_, inner_, af),
                    Dense(inner_, 1))
 
-init_params2 = Float64.(ComponentArray(Lux.setup(Random.default_rng(), chain2)[1]))
+initp, st = Lux.setup(Random.default_rng(), chain2)
+init_params2 = Float64.(ComponentArrays.ComponentArray(initp))
 
 @named pde_system = PDESystem(eq, bcs, domains, [x, y], [u(x, y)])
 
 losses = map(1:count_decomp) do i
-    loss(cord, θ) = chain2(cord, θ) .- phis[i](cord, reses[i].minimizer)
+    function loss(cord, θ)
+        global st, chain2, phis, reses
+        ch2, st = chain2(cord, θ, st)
+        ch2 .- phis[i](cord, reses[i].u)
+    end
 end
 
 callback = function (p, l)
@@ -244,15 +244,15 @@ end
 
 prob_ = NeuralPDE.neural_adapter(losses, init_params2, pde_system_map,
                                  NeuralPDE.QuadratureTraining())
-res_ = Optimization.solve(prob_, BFGS(); callback = callback, maxiters = 2000)
-prob_ = NeuralPDE.neural_adapter(losses, res_.minimizer, pde_system_map,
+res_ = Optimization.solve(prob_, OptimizationOptimisers.Adam(5e-3); maxiters = 5000)
+prob_ = NeuralPDE.neural_adapter(losses, res_.u, pde_system_map,
                                  NeuralPDE.QuadratureTraining())
-res_ = Optimization.solve(prob_, BFGS(); callback = callback, maxiters = 1000)
+res_ = Optimization.solve(prob_, OptimizationOptimisers.Adam(5e-3); maxiters = 5000)
 
-phi_ = NeuralPDE.get_phi(chain2)
+phi_ = PhysicsInformedNN(chain2, strategy; init_params = res_.u).phi
 
 xs, ys = [infimum(d.domain):dx:supremum(d.domain) for d in domains]
-u_predict_ = reshape([first(phi_([x, y], res_.minimizer)) for x in xs for y in ys],
+u_predict_ = reshape([first(phi_([x, y], res_.u)) for x in xs for y in ys],
                      (length(xs), length(ys)))
 u_real = reshape([analytic_sol_func(x, y) for x in xs for y in ys],
                  (length(xs), length(ys)))
@@ -267,5 +267,3 @@ p4 = plot(xs, ys, diff_u, linetype = :contourf, title = "error 1");
 p5 = plot(xs, ys, diff_u_, linetype = :contourf, title = "error 2");
 plot(p1, p2, p3, p4, p5)
 ```
-
-![decomp](https://user-images.githubusercontent.com/12683885/127651866-87ba95ad-10a8-471e-b4a5-093210a86c0e.png)
diff --git a/docs/src/tutorials/ode_parameter_estimation.md b/docs/src/tutorials/ode_parameter_estimation.md
index 2d3cba5517..39bd4ef229 100644
--- a/docs/src/tutorials/ode_parameter_estimation.md
+++ b/docs/src/tutorials/ode_parameter_estimation.md
@@ -10,6 +10,7 @@ We start by defining the problem:
 using NeuralPDE, OrdinaryDiffEq
 using Lux, Random
 using OptimizationOptimJL, LineSearches
+using Plots
 using Test # hide
 
 function lv(u, p, t)
@@ -40,7 +41,7 @@ Now, lets define a neural network for the PINN using [Lux.jl](https://lux.csail.
 ```@example param_estim_lv
 rng = Random.default_rng()
 Random.seed!(rng, 0)
-n = 12
+n = 15
 chain = Lux.Chain(
             Lux.Dense(1, n, Lux.σ),
             Lux.Dense(n, n, Lux.σ),
@@ -62,7 +63,7 @@ Next we define the optimizer and [`NNODE`](@ref) which is then plugged into the
 
 ```@example param_estim_lv
 opt = LBFGS(linesearch = BackTracking())
-alg = NNODE(chain, opt, ps; strategy = GridTraining(0.01), param_estim = true, additional_loss = additional_loss)
+alg = NNODE(chain, opt, ps; strategy = WeightedIntervalTraining([0.7, 0.2, 0.1], 500), param_estim = true, additional_loss = additional_loss)
 ```
 
 Now we have all the pieces to solve the optimization problem.
diff --git a/docs/src/tutorials/param_estim.md b/docs/src/tutorials/param_estim.md
index 2f2ebc782e..e696c76702 100644
--- a/docs/src/tutorials/param_estim.md
+++ b/docs/src/tutorials/param_estim.md
@@ -16,8 +16,8 @@ We start by defining the problem,
 
 ```@example param_estim
 using NeuralPDE, Lux, ModelingToolkit, Optimization, OptimizationOptimJL, OrdinaryDiffEq,
-      Plots
-import ModelingToolkit: Interval, infimum, supremum
+      Plots, LineSearches
+using ModelingToolkit: Interval, infimum, supremum
 @parameters t, σ_, β, ρ
 @variables x(..), y(..), z(..)
 Dt = Differential(t)
@@ -90,30 +90,11 @@ function additional_loss(phi, θ, p)
 end
 ```
 
-#### Note about Flux
-
-If Flux neural network is used, then the subsetting must be computed manually as `θ`
-is simply a vector. This looks like:
-
-```julia
-init_params = [Flux.destructure(c)[1] for c in [chain1, chain2, chain3]]
-acum = [0; accumulate(+, length.(init_params))]
-sep = [(acum[i] + 1):acum[i + 1] for i in 1:(length(acum) - 1)]
-(u_, t_) = data
-len = length(data[2])
-
-function additional_loss(phi, θ, p)
-    return sum(sum(abs2, phi[i](t_, θ[sep[i]]) .- u_[[i], :]) / len for i in 1:1:3)
-end
-```
-
-#### Back to our originally scheduled programming
-
 Then finally defining and optimizing using the `PhysicsInformedNN` interface.
 
 ```@example param_estim
 discretization = NeuralPDE.PhysicsInformedNN([chain1, chain2, chain3],
-                                             NeuralPDE.QuadratureTraining(), param_estim = true,
+                                             NeuralPDE.QuadratureTraining(; abstol = 1e-6, reltol = 1e-6, batch = 200), param_estim = true,
                                              additional_loss = additional_loss)
 @named pde_system = PDESystem(eqs, bcs, domains, [t], [x(t), y(t), z(t)], [σ_, ρ, β],
                               defaults = Dict([p .=> 1.0 for p in [σ_, ρ, β]]))
@@ -122,7 +103,7 @@ callback = function (p, l)
     println("Current loss is: $l")
     return false
 end
-res = Optimization.solve(prob, BFGS(); callback = callback, maxiters = 5000)
+res = Optimization.solve(prob, BFGS(linesearch = BackTracking()); maxiters = 1000)
 p_ = res.u[(end - 2):end] # p_ = [9.93, 28.002, 2.667]
 ```
 
@@ -135,5 +116,3 @@ u_predict = [[discretization.phi[i]([t], minimizers[i])[1] for t in ts] for i in
 plot(sol)
 plot!(ts, u_predict, label = ["x(t)" "y(t)" "z(t)"])
 ```
-
-![Plot_Lorenz](https://user-images.githubusercontent.com/12683885/110944192-2ae05f00-834d-11eb-910b-f5c06d22ec8a.png)
diff --git a/docs/src/tutorials/pdesystem.md b/docs/src/tutorials/pdesystem.md
index 4ace039bd4..d1c438b30c 100644
--- a/docs/src/tutorials/pdesystem.md
+++ b/docs/src/tutorials/pdesystem.md
@@ -27,9 +27,11 @@ Using physics-informed neural networks.
 
 ## Copy-Pasteable Code
 
-```@example
+```@example poisson
 using NeuralPDE, Lux, Optimization, OptimizationOptimJL
-import ModelingToolkit: Interval
+using LineSearches
+using ModelingToolkit: Interval
+using Plots
 
 @parameters x y
 @variables u(..)
@@ -51,26 +53,23 @@ dim = 2 # number of dimensions
 chain = Lux.Chain(Dense(dim, 16, Lux.σ), Dense(16, 16, Lux.σ), Dense(16, 1))
 
 # Discretization
-dx = 0.05
-discretization = PhysicsInformedNN(chain, QuadratureTraining())
+discretization = PhysicsInformedNN(chain, QuadratureTraining(; batch = 200, abstol = 1e-6, reltol = 1e-6))
 
 @named pde_system = PDESystem(eq, bcs, domains, [x, y], [u(x, y)])
 prob = discretize(pde_system, discretization)
 
-#Optimizer
-opt = OptimizationOptimJL.BFGS()
-
 #Callback function
 callback = function (p, l)
     println("Current loss is: $l")
     return false
 end
 
-res = Optimization.solve(prob, opt, callback = callback, maxiters = 1000)
+# Optimizer
+opt = OptimizationOptimJL.LBFGS(linesearch = BackTracking())
+res = solve(prob, opt, maxiters = 1000)
 phi = discretization.phi
 
-using Plots
-
+dx = 0.05
 xs, ys = [infimum(d.domain):(dx / 10):supremum(d.domain) for d in domains]
 analytic_sol_func(x, y) = (sin(pi * x) * sin(pi * y)) / (2pi^2)
 
@@ -92,7 +91,8 @@ The ModelingToolkit PDE interface for this example looks like this:
 
 ```@example poisson
 using NeuralPDE, Lux, ModelingToolkit, Optimization, OptimizationOptimJL
-import ModelingToolkit: Interval
+using ModelingToolkit: Interval
+using Plots
 
 @parameters x y
 @variables u(..)
@@ -122,8 +122,7 @@ Here, we build PhysicsInformedNN algorithm where `dx` is the step of discretizat
 `strategy` stores information for choosing a training strategy.
 
 ```@example poisson
-dx = 0.05
-discretization = PhysicsInformedNN(chain, QuadratureTraining())
+discretization = PhysicsInformedNN(chain, QuadratureTraining(; batch = 200, abstol = 1e-6, reltol = 1e-6))
 ```
 
 As described in the API docs, we now need to define the `PDESystem` and create PINNs
@@ -139,20 +138,22 @@ Here, we define the callback function and the optimizer. And now we can solve th
 
 ```@example poisson
 #Optimizer
-opt = OptimizationOptimJL.BFGS()
+opt = OptimizationOptimJL.LBFGS(linesearch = BackTracking())
 
 callback = function (p, l)
     println("Current loss is: $l")
     return false
 end
 
-res = Optimization.solve(prob, opt, callback = callback, maxiters = 1000)
+# We can pass the callback function in the solve. Not doing here as the output would be very long.
+res = Optimization.solve(prob, opt, maxiters = 1000)
 phi = discretization.phi
 ```
 
 We can plot the predicted solution of the PDE and compare it with the analytical solution to plot the relative error.
 
 ```@example poisson
+dx = 0.05
 xs, ys = [infimum(d.domain):(dx / 10):supremum(d.domain) for d in domains]
 analytic_sol_func(x, y) = (sin(pi * x) * sin(pi * y)) / (2pi^2)
 
@@ -162,12 +163,8 @@ u_real = reshape([analytic_sol_func(x, y) for x in xs for y in ys],
                  (length(xs), length(ys)))
 diff_u = abs.(u_predict .- u_real)
 
-using Plots
-
 p1 = plot(xs, ys, u_real, linetype = :contourf, title = "analytic");
 p2 = plot(xs, ys, u_predict, linetype = :contourf, title = "predict");
 p3 = plot(xs, ys, diff_u, linetype = :contourf, title = "error");
 plot(p1, p2, p3)
 ```
-
-![poissonplot](https://user-images.githubusercontent.com/12683885/90962648-2db35980-e4ba-11ea-8e58-f4f07c77bcb9.png)
diff --git a/docs/src/tutorials/systems.md b/docs/src/tutorials/systems.md
index 9e20bc7f73..80e3fbb62a 100644
--- a/docs/src/tutorials/systems.md
+++ b/docs/src/tutorials/systems.md
@@ -35,8 +35,8 @@ with physics-informed neural networks.
 ## Solution
 
 ```@example system
-using NeuralPDE, Lux, ModelingToolkit, Optimization, OptimizationOptimJL
-import ModelingToolkit: Interval, infimum, supremum
+using NeuralPDE, Lux, ModelingToolkit, Optimization, OptimizationOptimJL, LineSearches
+using ModelingToolkit: Interval, infimum, supremum
 
 @parameters t, x
 @variables u1(..), u2(..), u3(..)
@@ -67,7 +67,7 @@ input_ = length(domains)
 n = 15
 chain = [Lux.Chain(Dense(input_, n, Lux.σ), Dense(n, n, Lux.σ), Dense(n, 1)) for _ in 1:3]
 
-strategy = QuadratureTraining()
+strategy = QuadratureTraining(; batch = 200, abstol = 1e-6, reltol = 1e-6)
 discretization = PhysicsInformedNN(chain, strategy)
 
 @named pdesystem = PDESystem(eqs, bcs, domains, [t, x], [u1(t, x), u2(t, x), u3(t, x)])
@@ -79,13 +79,12 @@ bcs_inner_loss_functions = sym_prob.loss_functions.bc_loss_functions
 
 callback = function (p, l)
     println("loss: ", l)
-    println("pde_losses: ", map(l_ -> l_(p), pde_inner_loss_functions))
-    println("bcs_losses: ", map(l_ -> l_(p), bcs_inner_loss_functions))
+    println("pde_losses: ", map(l_ -> l_(p.u), pde_inner_loss_functions))
+    println("bcs_losses: ", map(l_ -> l_(p.u), bcs_inner_loss_functions))
     return false
 end
 
-res = Optimization.solve(prob, BFGS(); callback = callback, maxiters = 5000)
-
+res = solve(prob, LBFGS(linesearch = BackTracking()); maxiters = 1000)
 phi = discretization.phi
 ```
 
@@ -96,7 +95,7 @@ interface. Here is an example using the components from `symbolic_discretize` to
 reproduce the `discretize` optimization:
 
 ```@example system
-using NeuralPDE, Lux, ModelingToolkit, Optimization, OptimizationOptimJL
+using NeuralPDE, Lux, ModelingToolkit, Optimization, OptimizationOptimJL, LineSearches
 import ModelingToolkit: Interval, infimum, supremum
 
 @parameters t, x
@@ -138,8 +137,8 @@ bc_loss_functions = sym_prob.loss_functions.bc_loss_functions
 
 callback = function (p, l)
     println("loss: ", l)
-    println("pde_losses: ", map(l_ -> l_(p), pde_loss_functions))
-    println("bcs_losses: ", map(l_ -> l_(p), bc_loss_functions))
+    println("pde_losses: ", map(l_ -> l_(p.u), pde_loss_functions))
+    println("bcs_losses: ", map(l_ -> l_(p.u), bc_loss_functions))
     return false
 end
 
@@ -152,8 +151,7 @@ end
 f_ = OptimizationFunction(loss_function, Optimization.AutoZygote())
 prob = Optimization.OptimizationProblem(f_, sym_prob.flat_init_params)
 
-res = Optimization.solve(prob, OptimizationOptimJL.BFGS(); callback = callback,
-                         maxiters = 5000)
+res = Optimization.solve(prob, OptimizationOptimJL.LBFGS(linesearch = BackTracking()); maxiters = 1000)
 ```
 
 ## Solution Representation
@@ -174,20 +172,26 @@ end
 u_real = [[analytic_sol_func(t, x)[i] for t in ts for x in xs] for i in 1:3]
 u_predict = [[phi[i]([t, x], minimizers_[i])[1] for t in ts for x in xs] for i in 1:3]
 diff_u = [abs.(u_real[i] .- u_predict[i]) for i in 1:3]
+ps = []
 for i in 1:3
     p1 = plot(ts, xs, u_real[i], linetype = :contourf, title = "u$i, analytic")
     p2 = plot(ts, xs, u_predict[i], linetype = :contourf, title = "predict")
     p3 = plot(ts, xs, diff_u[i], linetype = :contourf, title = "error")
-    plot(p1, p2, p3)
-    savefig("sol_u$i")
+    push!(ps, plot(p1, p2, p3))
 end
 ```
 
-![sol_uq1](https://user-images.githubusercontent.com/12683885/122979254-03634e80-d3a0-11eb-985b-d3bae2dddfde.png)
+```@example system
+ps[1]
+```
 
-![sol_uq2](https://user-images.githubusercontent.com/12683885/122979278-09592f80-d3a0-11eb-8fee-de3652f138d8.png)
+```@example system
+ps[2]
+```
 
-![sol_uq3](https://user-images.githubusercontent.com/12683885/122979288-0e1de380-d3a0-11eb-9005-bfb501959b83.png)
+```@example system
+ps[3]
+```
 
 Notice here that the solution is represented in the `OptimizationSolution` with `u` as
 the parameters for the trained neural network. But, for the case where the neural network
@@ -200,25 +204,16 @@ Subsetting the array also works, but is inelegant.
 
 (If `param_estim == true`, then `res.u.p` are the fit parameters)
 
-If Flux.jl is used, then subsetting the array is required. This looks like:
-
-```julia
-init_params = [Flux.destructure(c)[1] for c in chain]
-acum = [0; accumulate(+, length.(init_params))]
-sep = [(acum[i] + 1):acum[i + 1] for i in 1:(length(acum) - 1)]
-minimizers_ = [res.minimizer[s] for s in sep]
-```
-
 #### Note: Solving Matrices of PDEs
 
 Also, in addition to vector systems, we can use the matrix form of PDEs:
 
-```@example
+```julia
 using ModelingToolkit, NeuralPDE
 @parameters x y
-@variables u[1:2, 1:2](..)
-@derivatives Dxx'' ~ x
-@derivatives Dyy'' ~ y
+@variables (u(..))[1:2, 1:2]
+Dxx = Differential(x)^2
+Dyy = Differential(y)^2
 
 # Initial and boundary conditions
 bcs = [u[1](x, 0) ~ x, u[2](x, 0) ~ 2, u[3](x, 0) ~ 3, u[4](x, 0) ~ 4]