diff --git a/docs/Project.toml b/docs/Project.toml
index c6d1296..4ef18fb 100644
--- a/docs/Project.toml
+++ b/docs/Project.toml
@@ -2,6 +2,8 @@
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
 HighDimPDE = "57c578d5-59d4-4db8-a490-a9fc372d19d2"
+LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+StochasticDiffEq = "789caeaf-c7a9-5a7d-9973-96adeb23e2a0"
 
 [compat]
 Documenter = "1"
diff --git a/docs/make.jl b/docs/make.jl
index 896ce68..ea0bd67 100644
--- a/docs/make.jl
+++ b/docs/make.jl
@@ -6,7 +6,7 @@ cp("./docs/Project.toml", "./docs/src/assets/Project.toml", force = true)
 include("pages.jl")
 
 makedocs(sitename = "HighDimPDE.jl",
-    authors = "Victor Boussange",
+    authors = "#",
     pages = pages,
     clean = true, doctest = false, linkcheck = true,
     format = Documenter.HTML(assets = ["assets/favicon.ico"],
diff --git a/docs/src/NNKolmogorov.md b/docs/src/NNKolmogorov.md
index 73d554b..669f01a 100644
--- a/docs/src/NNKolmogorov.md
+++ b/docs/src/NNKolmogorov.md
@@ -1,4 +1,4 @@
-# [The `NNKolmogorov` algorithm](@id nn_komogorov)
+# [The `NNKolmogorov` algorithm](@id nn_kolmogorov)
 
 ### Problems Supported:
 1. [`ParabolicPDEProblem`](@ref)
diff --git a/docs/src/NNParamKolmogorov.md b/docs/src/NNParamKolmogorov.md
index 5635f8c..0b36d93 100644
--- a/docs/src/NNParamKolmogorov.md
+++ b/docs/src/NNParamKolmogorov.md
@@ -1,4 +1,4 @@
-# [The `NNParamKolmogorov` algorithm](@id nn_komogorov)
+# [The `NNParamKolmogorov` algorithm](@id nn_paramkolmogorov)
 
 ### Problems Supported:
 1. [`ParabolicPDEProblem`](@ref)
diff --git a/docs/src/assets/Project.toml b/docs/src/assets/Project.toml
new file mode 100644
index 0000000..4ef18fb
--- /dev/null
+++ b/docs/src/assets/Project.toml
@@ -0,0 +1,11 @@
+[deps]
+Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
+Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
+HighDimPDE = "57c578d5-59d4-4db8-a490-a9fc372d19d2"
+LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+StochasticDiffEq = "789caeaf-c7a9-5a7d-9973-96adeb23e2a0"
+
+[compat]
+Documenter = "1"
+Flux = "0.13, 0.14"
+HighDimPDE = "2"
diff --git a/docs/src/getting_started.md b/docs/src/getting_started.md
index dfaff3c..f7da87f 100644
--- a/docs/src/getting_started.md
+++ b/docs/src/getting_started.md
@@ -32,8 +32,8 @@ x0 = fill(0., d)  # initial point
 g(x) = exp(-sum(x.^2)) # initial condition
 μ(x, p, t) = 0.0 # advection coefficients
 σ(x, p, t) = 0.1 # diffusion coefficients
-f(x, y, v_x, v_y, ∇v_x, ∇v_y, p, t) = max(0.0, v_x) * (1 -  max(0.0, v_x)) # nonlocal nonlinear part of the
-prob = PIDEProblem(μ, σ, x0, tspan, g, f) # defining the problem
+f(x, v_x, ∇v_x, p, t) = max(0.0, v_x) * (1 -  max(0.0, v_x)) # nonlocal nonlinear part of the
+prob = ParabolicPDEProblem(μ, σ, x0, tspan; g, f) # defining the problem
 
 ## Definition of the algorithm
 alg = MLP() # defining the algorithm. We use the Multi Level Picard algorithm
@@ -117,7 +117,7 @@ sol = solve(prob,
 
 [`DeepSplitting`](@ref deepsplitting) can run on the GPU for (much) improved performance. To do so, just set `use_cuda = true`.
 
-```@example DeepSplitting_gpu
+```@example DeepSplitting_non_local_PDE
 sol = solve(prob, 
             alg, 
             0.1, 
diff --git a/docs/src/index.md b/docs/src/index.md
index 32c70e0..7c50a08 100644
--- a/docs/src/index.md
+++ b/docs/src/index.md
@@ -3,7 +3,7 @@
 
 **HighDimPDE.jl** is a Julia package to **solve Highly Dimensional non-linear, non-local PDEs** of the forms: 
 
-1. Partial Integro Differential Equations: 
+### 1. Partial Integro Differential Equations: 
 ```math
 \begin{aligned}
     (\partial_t u)(t,x) &= \int_{\Omega} f\big(t,x,{\bf x}, u(t,x),u(t,{\bf x}), ( \nabla_x u )(t,x ),( \nabla_x u )(t,{\bf x} ) \big) \, d{\bf x} \\
@@ -11,9 +11,9 @@
 \end{aligned}
 ```
 
-where $u \colon [0,T] \times \Omega \to \R$, $\Omega \subseteq \R^d$ is subject to initial and boundary conditions, and where $d$ is large.
+where $u \colon [0,T] \times \Omega \to \R$, $\Omega \subseteq \R^d$ is subject to initial and boundary conditions, and where $d$ is large. These equations are defined using the [`PIDEProblem`](@ref)
 
-2. Parabolic Partial Differential Equations: 
+### 2. Parabolic Partial Differential Equations: 
 ```math
 \begin{aligned}
     (\partial_t u)(t,x) &=  f\big(t,x, u(t,x), ( \nabla_x u )(t,x )\big) 
@@ -21,7 +21,7 @@ where $u \colon [0,T] \times \Omega \to \R$, $\Omega \subseteq \R^d$ is subject
 \end{aligned}
 ```
 
-where $u \colon [0,T] \times \Omega \to \R$, $\Omega \subseteq \R^d$ is subject to initial and boundary conditions, and where $d$ is large.
+where $u \colon [0,T] \times \Omega \to \R$, $\Omega \subseteq \R^d$ is subject to initial and boundary conditions, and where $d$ is large. These equations are defined using the [`ParabolicPDEProblem`](@ref)
 
 !!! note
     The difference between the two problems is that in Partial Integro Differential Equations, the integrand is integrated over **x**, while in Parabolic Integro Differential Equations, the function `f` is just evaluated for `x`.
@@ -32,8 +32,11 @@ where $u \colon [0,T] \times \Omega \to \R$, $\Omega \subseteq \R^d$ is subject
 
 * the [Multi-Level Picard iterations scheme](@ref mlp)
 
-* the Deep BSDE scheme (@ref deepbsde).
+* [the Deep BSDE scheme](@ref deepbsde).
 
+* [NNKolmogorov](@ref nn_kolmogorov) and [NNParamKolmogorov](@ref nn_paramkolmogorov) schemes for Kolmogorov PDEs.
+
+* [NNStopping] scheme for solving optimal stopping time problem.
 
 To make the most out of **HighDimPDE.jl**, we advise to first have a look at the 
 
diff --git a/docs/src/tutorials/deepbsde.md b/docs/src/tutorials/deepbsde.md
index 1e11fa5..f439e18 100644
--- a/docs/src/tutorials/deepbsde.md
+++ b/docs/src/tutorials/deepbsde.md
@@ -5,33 +5,46 @@ Hamilton-Jacobi-Bellman equation that takes a few minutes on a laptop:
 
 ```@example deepbsde
 using HighDimPDE
-using Flux, OptimizationOptimisers
-using DifferentialEquations
+using Flux
+using StochasticDiffEq
 using LinearAlgebra
+
 d = 100 # number of dimensions
-X0 = fill(0.0f0, d) # initial value of stochastic control process
+x0 = fill(0.0f0, d)
 tspan = (0.0f0, 1.0f0)
+dt = 0.2f0
 λ = 1.0f0
+#
+g(X) = log(0.5f0 + 0.5f0 * sum(X .^ 2))
+f(X, u, σᵀ∇u, p, t) = -λ * sum(σᵀ∇u .^ 2)
+μ_f(X, p, t) = zero(X)  #Vector d x 1 λ
+σ_f(X, p, t) = Diagonal(sqrt(2.0f0) * ones(Float32, d)) #Matrix d x d
+prob = ParabolicPDEProblem(μ_f, σ_f, x0, tspan; g, f)
 
-g(X) = log(0.5f0 + 0.5f0 * sum(X.^2))
-f(X,u,σᵀ∇u,p,t) = -λ * sum(σᵀ∇u.^2)
-μ_f(X,p,t) = zero(X)  # Vector d x 1 λ
-σ_f(X,p,t) = Diagonal(sqrt(2.0f0) * ones(Float32, d)) # Matrix d x d
-prob = PIDEProblem(μ_f, σ_f, X0, tspan, g, f)
-hls = 10 + d # hidden layer size
-opt = Optimisers.Adam(0.01)  # optimizer
-# sub-neural network approximating solutions at the desired point
+hls = 10 + d #hidden layer size
+opt = Flux.Optimise.Adam(0.1)  #optimizer
+#sub-neural network approximating solutions at the desired point
 u0 = Flux.Chain(Dense(d, hls, relu),
-                Dense(hls, hls, relu),
-                Dense(hls, 1))
+    Dense(hls, hls, relu),
+    Dense(hls, hls, relu),
+    Dense(hls, 1))
 # sub-neural network approximating the spatial gradients at time point
 σᵀ∇u = Flux.Chain(Dense(d + 1, hls, relu),
-                  Dense(hls, hls, relu),
-                  Dense(hls, hls, relu),
-                  Dense(hls, d))
-pdealg = NNPDENS(u0, σᵀ∇u, opt=opt)
-@time ans = solve(prob, pdealg, verbose=true, maxiters=100, trajectories=100,
-                            alg=EM(), dt=1.2, pabstol=1f-2)
+    Dense(hls, hls, relu),
+    Dense(hls, hls, relu),
+    Dense(hls, hls, relu),
+    Dense(hls, d))
+pdealg = DeepBSDE(u0, σᵀ∇u, opt = opt)
+
+@time sol = solve(prob,
+    pdealg,
+    StochasticDiffEq.EM(),
+    verbose = true,
+    maxiters = 150,
+    trajectories = 30,
+    dt = 1.2f0,
+    pabstol = 1.0f-4)
+
 ```
 
 Now, let's explain the details!
@@ -61,9 +74,14 @@ with terminating condition $g(x) = \log(1/2 + 1/2 \|x\|^2))$.
 
 #### Define the Problem
 
-To get the solution above using the `PIDEProblem`, we write:
+To get the solution above using the [`ParabolicPDEProblem`](@ref), we write:
 
 ```@example deepbsde2
+using HighDimPDE
+using Flux
+using StochasticDiffEq
+using LinearAlgebra
+
 d = 100 # number of dimensions
 X0 = fill(0.0f0,d) # initial value of stochastic control process
 tspan = (0.0f0, 1.0f0)
@@ -73,12 +91,12 @@ g(X) = log(0.5f0 + 0.5f0*sum(X.^2))
 f(X,u,σᵀ∇u,p,t) = -λ*sum(σᵀ∇u.^2)
 μ_f(X,p,t) = zero(X)  #Vector d x 1 λ
 σ_f(X,p,t) = Diagonal(sqrt(2.0f0)*ones(Float32,d)) #Matrix d x d
-prob = PIDEProblem(μ_f, σ_f, X0, tspan, g, f)
+prob = ParabolicPDEProblem(μ_f, σ_f, X0, tspan; g, f)
 ```
 
 #### Define the Solver Algorithm
 
-As described in the API docs, we now need to define our `NNPDENS` algorithm
+As described in the API docs, we now need to define our `DeepBSDE` algorithm
 by giving it the Flux.jl chains we want it to use for the neural networks.
 `u0` needs to be a `d` dimensional -> 1 dimensional chain, while `σᵀ∇u`
 needs to be `d+1` dimensional to `d` dimensions. Thus we define the following:
@@ -95,14 +113,13 @@ u0 = Flux.Chain(Dense(d,hls,relu),
                   Dense(hls,hls,relu),
                   Dense(hls,hls,relu),
                   Dense(hls,d))
-pdealg = NNPDENS(u0, σᵀ∇u, opt=opt)
+pdealg = DeepBSDE(u0, σᵀ∇u, opt = opt)
 ```
 
 #### Solving with Neural Nets
 
 ```@example deepbsde2
-@time ans = solve(prob, pdealg, verbose=true, maxiters=100, trajectories=100,
-                            alg=EM(), dt=0.2, pabstol = 1f-2)
+@time ans = solve(prob, pdealg, EM(), verbose=true, maxiters=100, trajectories=100, dt=0.2f0, pabstol = 1f-2)
 
 ```
 
@@ -121,7 +138,7 @@ equation, which models uncertain volatility and interest rates derived from the
 Black-Scholes equation. This model results in a nonlinear PDE whose dimension
 is the number of assets in the portfolio.
 
-To solve it using the `PIDEProblem`, we write:
+To solve it using the `ParabolicPDEProblem`, we write:
 
 ```julia
 d = 100 # number of dimensions
@@ -133,7 +150,7 @@ f(X,u,σᵀ∇u,p,t) = r * (u - sum(X.*σᵀ∇u))
 g(X) = sum(X.^2)
 μ_f(X,p,t) = zero(X) #Vector d x 1
 σ_f(X,p,t) = Diagonal(sigma*X) #Matrix d x d
-prob = PIDEProblem(μ_f, σ_f, X0, tspan, g, f)
+prob = ParabolicPDEProblem(μ_f, σ_f, X0, tspan; g, f)
 ```
 
 As described in the API docs, we now need to define our `NNPDENS` algorithm
@@ -151,7 +168,7 @@ u0 = Flux.Chain(Dense(d,hls,relu),
                   Dense(hls,hls,relu),
                   Dense(hls,hls,relu),
                   Dense(hls,d))
-pdealg = NNPDENS(u0, σᵀ∇u, opt=opt)
+pdealg = DeepBSDE(u0, σᵀ∇u, opt=opt)
 ```
 
 And now we solve the PDE. Here, we say we want to solve the underlying neural
@@ -160,8 +177,7 @@ SDE using the Euler-Maruyama SDE solver with our chosen `dt=0.2`, do at most
 and stop if the loss ever goes below `1f-6`.
 
 ```julia
-ans = solve(prob, pdealg, verbose=true, maxiters=150, trajectories=100,
-                            alg=EM(), dt=0.2, pabstol = 1f-6)
+ans = solve(prob, pdealg, EM(), verbose=true, maxiters=150, trajectories=100, dt=0.2f0)
 ```
 
 ## References
diff --git a/docs/src/tutorials/deepsplitting.md b/docs/src/tutorials/deepsplitting.md
index 11f47ea..138c7c9 100644
--- a/docs/src/tutorials/deepsplitting.md
+++ b/docs/src/tutorials/deepsplitting.md
@@ -11,7 +11,7 @@ where $\Omega = [-1/2, 1/2]^d$, and let's assume Neumann Boundary condition on $
 Let's solve Eq. (1) with the [`DeepSplitting`](@ref deepsplitting) solver. 
 
 ```@example deepsplitting
-using HighDimPDE
+using HighDimPDE, Flux
 
 ## Definition of the problem
 d = 10 # dimension of the problem
@@ -22,10 +22,13 @@ g(x) = exp.(- sum(x.^2, dims=1) ) # initial condition
 σ(x, p, t) = 0.1f0 # diffusion coefficients
 x0_sample = UniformSampling(fill(-5f-1, d), fill(5f-1, d))
 f(x, y, v_x, v_y, ∇v_x, ∇v_y, p, t) = v_x .* (1f0 .- v_y)
+```
+Since this is a non-local equation, we will define our problem as a [`PIDEProblem`](@ref)
+```@example deepsplitting
 prob = PIDEProblem(μ, σ, x0, tspan, g, f; x0_sample = x0_sample)
-
+```
+```@example deepsplitting
 ## Definition of the neural network to use
-using Flux # needed to define the neural network
 
 hls = d + 50 #hidden layer size
 
@@ -53,7 +56,7 @@ sol = solve(prob,
 
 `DeepSplitting` can run on the GPU for (much) improved performance. To do so, just set `use_cuda = true`.
 
-```@example deepsplitting2
+```@example deepsplitting
 sol = solve(prob, 
             alg, 
             0.1, 
@@ -62,4 +65,50 @@ sol = solve(prob,
             maxiters = 1000,
             batch_size = 1000,
             use_cuda=true)
+```
+
+# Solving local Allen Cahn Equation with Neumann BC
+```@example deepsplitting2 
+using HighDimPDE, Flux, StochasticDiffEq
+batch_size = 1000
+train_steps = 1000
+
+tspan = (0.0f0, 5.0f-1)
+dt = 5.0f-2  # time step
+
+μ(x, p, t) = 0.0f0 # advection coefficients
+σ(x, p, t) = 1.0f0 # diffusion coefficients
+
+d = 10
+∂ = fill(5.0f-1, d)
+
+hls = d + 50 #hidden layer size
+
+nn = Flux.Chain(Dense(d, hls, relu),
+        Dense(hls, hls, relu),
+        Dense(hls, 1)) # Neural network used by the scheme
+
+opt = Flux.Optimise.Adam(1e-2) #optimiser
+alg = DeepSplitting(nn, opt = opt)
+
+X0 = fill(0.0f0, d)  # initial point
+g(X) = exp.(-0.25f0 * sum(X .^ 2, dims = 1))   # initial condition
+a(u) = u - u^3
+f(y, v_y, ∇v_y, p, t) = a.(v_y) # nonlocal nonlinear function
+```
+Since we are dealing with a local problem here, i.e. no integration term used, we use [`ParabolicPDEProblem`](@ref) to define the problem.
+```@example deepsplitting2
+# defining the problem
+prob = ParabolicPDEProblem(μ, σ, X0, tspan; g, f, neumann_bc = [-∂, ∂])
+```
+```@example deepsplitting2
+# solving
+@time sol = solve(prob,
+        alg,
+        dt,
+        verbose = false,
+        abstol = 1e-5,
+        use_cuda = false,
+        maxiters = train_steps,
+        batch_size = batch_size)
 ```
\ No newline at end of file
diff --git a/docs/src/tutorials/mlp.md b/docs/src/tutorials/mlp.md
index 4756949..15ecf69 100644
--- a/docs/src/tutorials/mlp.md
+++ b/docs/src/tutorials/mlp.md
@@ -17,7 +17,7 @@ g(x) = exp(- sum(x.^2) ) # initial condition
 μ(x, p, t) = 0.0 # advection coefficients
 σ(x, p, t) = 0.1 # diffusion coefficients
 f(x, v_x, ∇v_x, p, t) = max(0.0, v_x) * (1 -  max(0.0, v_x)) # nonlocal nonlinear part of the
-prob = ParabolicPDEProblem(μ, σ, x0, tspan, g, f) # defining the problem
+prob = ParabolicPDEProblem(μ, σ, x0, tspan; g, f) # defining the problem
 
 ## Definition of the algorithm
 alg = MLP() # defining the algorithm. We use the Multi Level Picard algorithm
@@ -43,7 +43,7 @@ g(x) = exp( -sum(x.^2) ) # initial condition
 μ(x, p, t) = 0.0 # advection coefficients
 σ(x, p, t) = 0.1 # diffusion coefficients
 mc_sample = UniformSampling(fill(-5f-1, d), fill(5f-1, d))
-f(x, y, v_x, v_y, ∇v_x, ∇v_y, t) = max(0.0, v_x) * (1 -  max(0.0, v_y)) 
+f(x, y, v_x, v_y, ∇v_x, ∇v_y, p, t) = max(0.0, v_x) * (1 -  max(0.0, v_y)) 
 prob = PIDEProblem(μ, σ, x0, tspan, g, f) # defining x0_sample is sufficient to implement Neumann boundary conditions
 
 ## Definition of the algorithm
diff --git a/docs/src/tutorials/nnkolmogorov.md b/docs/src/tutorials/nnkolmogorov.md
index 9f35f95..29fd355 100644
--- a/docs/src/tutorials/nnkolmogorov.md
+++ b/docs/src/tutorials/nnkolmogorov.md
@@ -2,7 +2,9 @@
 
 ## Solving high dimensional Rainbow European Options for a range of initial stock prices:
 
-```julia
+```@example nnkolmogorov
+using HighDimPDE, StochasticDiffEq, Flux, LinearAlgebra
+
 d = 10 # dims
 T = 1/12 
 sigma = 0.01 .+ 0.03.*Matrix(Diagonal(ones(d))) # volatility 
@@ -22,12 +24,12 @@ xspan = [(98.00, 102.00) for i in 1:d]
 
 g(x) = max(maximum(x) -K, 0)
 
-sdealg = EM()
+sdealg = SRIW1()
 # provide `x0` as nothing to the problem since we are provinding a range for `x0`.
 prob = ParabolicPDEProblem(μ_func, σ_func, nothing, tspan, g = g, xspan = xspan)
 opt = Flux.Optimisers.Adam(0.01)
-alg = NNKolmogorov(m, opt)
 m = Chain(Dense(d, 16, elu), Dense(16, 32, elu), Dense(32, 16, elu), Dense(16, 1))
+alg = NNKolmogorov(m, opt)
 sol = solve(prob, alg, sdealg, verbose = true, dt = 0.01,
     dx = 0.0001, trajectories = 1000, abstol = 1e-6, maxiters = 300)
 ```
diff --git a/docs/src/tutorials/nnparamkolmogorov.md b/docs/src/tutorials/nnparamkolmogorov.md
index 735cf2d..f9de711 100644
--- a/docs/src/tutorials/nnparamkolmogorov.md
+++ b/docs/src/tutorials/nnparamkolmogorov.md
@@ -3,7 +3,8 @@
 ## Solving Parametric Family of High Dimensional Heat Equation.
 
 In this example we will solve the high dimensional heat equation over a range of initial values, and also over a range of thermal diffusivity.
-```julia
+```@example nnparamkolmogorov
+using HighDimPDE, Flux, StochasticDiffEq
 d = 10
 # models input is `d` for initial values, `d` for thermal diffusivity, and last dimension is for stopping time.
 m = Chain(Dense(d + 1 + 1, 32, relu), Dense(32, 16, relu), Dense(16, 8, relu), Dense(8, 1))
@@ -49,13 +50,13 @@ sol = solve(prob, NNParamKolmogorov(m, opt), sdealg, verbose = true, dt = 0.01,
 Similarly we can parametrize the drift function `mu_` and the initial function `g`, and obtain a solution over all parameters and initial values.
 
 # Inferring on the solution from `NNParamKolmogorov`:
-```julia
+```@example nnparamkolmogorov
 x_test = rand(xspan[1][1]:0.1:xspan[1][2], d)
 p_sigma_test = rand(p_domain.p_sigma[1]:dps.p_sigma:p_domain.p_sigma[2], 1, 1)
 t_test = rand(tspan[1]:dt:tspan[2], 1, 1)
 p_mu_test = nothing
 p_phi_test = nothing
 ```
-```julia
+```@example nnparamkolmogorov
 sol.ufuns(x_test, t_test, p_sigma_test, p_mu_test, p_phi_test)
 ```
\ No newline at end of file
diff --git a/docs/src/tutorials/nnstopping.md b/docs/src/tutorials/nnstopping.md
index 25edad2..486fa8e 100644
--- a/docs/src/tutorials/nnstopping.md
+++ b/docs/src/tutorials/nnstopping.md
@@ -8,11 +8,12 @@ We will calculate optimal strategy for Bermudan Max-Call option with following d
 g(t, x) =  e^{-rt}max\lbrace max\lbrace x1, ... , xd \rbrace − K, 0\rbrace
 ```
 We define the parameters, drift function and the diffusion function for the dynamics of the option.
-```julia
+```@example nnstopping
+using HighDimPDE, Flux, StochasticDiffEq
 d = 3 # Number of assets in the stock
 r = 0.05 # interest rate
 beta = 0.2 # volatility
-T = 3 # maturity
+T = 3.0 # maturity
 u0 = fill(90.0, d) # initial stock value
 delta = 0.1 # delta
 f(du, u, p, t) = du .= (r - delta) * u # drift
@@ -28,15 +29,16 @@ function g(x, t)
 end
 
 ```
-We then define a `PIDEProblem` with no non linear term:
-```julia
-prob = PIDEProblem(f, sigma, u0, tspan; payoff = g)
+We then define a [`ParabolicPDEProblem`](@ref) with no non linear term:
+
+```@example nnstopping
+prob = ParabolicPDEProblem(f, sigma, u0, tspan; payoff = g)
 ```
 !!! note 
     We provide the payoff function with a keyword argument `payoff` 
 
 And now we define our models:
-```julia
+```@example nnstopping
 models = [Chain(Dense(d + 1, 32, tanh), BatchNorm(32, tanh), Dense(32, 1, sigmoid))
           for i in 1:N]
 ```
@@ -44,7 +46,7 @@ models = [Chain(Dense(d + 1, 32, tanh), BatchNorm(32, tanh), Dense(32, 1, sigmoi
     The number of models should be equal to the time discritization.
 
 And finally we define our optimizer and algorithm, and call `solve`:
-```julia
+```@example nnstopping
 opt = Flux.Optimisers.Adam(0.01)
 alg = NNStopping(models, opt)
 
diff --git a/src/DeepSplitting.jl b/src/DeepSplitting.jl
index 91ca483..0722ee9 100644
--- a/src/DeepSplitting.jl
+++ b/src/DeepSplitting.jl
@@ -10,7 +10,7 @@ end
 Deep splitting algorithm.
 
 # Arguments
-* `nn`: a [Flux.Chain](https://fluxml.ai/Flux.jl/stable/models/layers/#Flux.Chain), or more generally a [functor](https://github.com/FluxML/Functors.jl).
+* `nn`: a [Flux.Chain](https://fluxml.ai/Flux.jl/stable/reference/models/layers/#Flux.Chain), or more generally a [functor](https://github.com/FluxML/Functors.jl).
 * `K`: the number of Monte Carlo integrations.
 * `opt`: optimizer to be used. By default, `Flux.Optimise.Adam(0.01)`.
 * `λs`: the learning rates, used sequentially. Defaults to a single value taken from `opt`.