diff --git a/Project.toml b/Project.toml
index f951e1313..f0dfcbf8c 100644
--- a/Project.toml
+++ b/Project.toml
@@ -1,7 +1,7 @@
 name = "JUDI"
 uuid = "f3b833dc-6b2e-5b9c-b940-873ed6319979"
 authors = ["Philipp Witte, Mathias Louboutin"]
-version = "3.4.0"
+version = "3.4.1"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
diff --git a/docker/Dockerfile.JUDI b/docker/Dockerfile.JUDI
index 7d329b11c..7549947e9 100644
--- a/docker/Dockerfile.JUDI
+++ b/docker/Dockerfile.JUDI
@@ -1,6 +1,9 @@
 ARG JVER=1.7
 FROM mloubout/julia-base:$JVER
 
+# Update devito to latest
+RUN pip3 install --user --upgrade devito[extras,tests]
+
 # Update julia packages
 RUN julia -e 'using Pkg;Pkg.update()'
 
diff --git a/docs/src/inversion.md b/docs/src/inversion.md
index cb9661a91..0217efd39 100644
--- a/docs/src/inversion.md
+++ b/docs/src/inversion.md
@@ -92,8 +92,8 @@ res = spg(objective_function, vec(m0), ProjBound, options)
 This example script can be run in parallel and requires roughly 220 MB of memory per source location. Execute the following code to generate figures of the initial model and the result, as well as the function values:
 
 ```julia
-figure(); imshow(sqrt.(1./adjoint(m0))); title("Initial model")
-figure(); imshow(sqrt.(1./adjoint(reshape(x, model0.n)))); title("FWI")
+figure(); imshow(sqrt.(1. /adjoint(m0))); title("Initial model")
+figure(); imshow(sqrt.(1. /adjoint(reshape(x, model0.n)))); title("FWI")
 figure(); plot(fvals); title("Function value")
 ```
 
diff --git a/docs/src/preconditioners.md b/docs/src/preconditioners.md
index 11eab70ad..a02937579 100644
--- a/docs/src/preconditioners.md
+++ b/docs/src/preconditioners.md
@@ -129,5 +129,5 @@ lsrtm_objective(model, q, dobs, dm; model_precon=dPrec, data_precon=dmPrec)
 where `dPrec` and `dmPrec` can be:
 
 - A single preconditioner ([DataPreconditionner](@ref) for `data_precon` and [ModelPreconditionner](@ref) for `model_precon`)
-- A list/tuple of preconditioners ([DataPreconditionner](@ref) for `data_precon` and ModelPreconditionner](@ref) for `model_precon`)
-- A product of preconditioners ([DataPreconditionner](@ref) for `data_precon` and ModelPreconditionner](@ref) for `model_precon`)
+- A list/tuple of preconditioners ([DataPreconditionner](@ref) for `data_precon` and [ModelPreconditionner](@ref) for `model_precon`)
+- A product of preconditioners ([DataPreconditionner](@ref) for `data_precon` and [ModelPreconditionner](@ref) for `model_precon`)
diff --git a/src/TimeModeling/Modeling/propagation.jl b/src/TimeModeling/Modeling/propagation.jl
index 142b6e960..8283ce305 100644
--- a/src/TimeModeling/Modeling/propagation.jl
+++ b/src/TimeModeling/Modeling/propagation.jl
@@ -102,7 +102,7 @@ This is the main multi-source wrapper function for `fwi_objective` and `lsrtm_ob
 Computes the misifit and gradient (LSRTM if `lin` else FWI) for the given `q` source and `dobs` and
 perturbation `dm`.
 """
-function multi_src_fg!(G, model, q, dobs, dm; options=Options(), kw...)
+function multi_src_fg!(G, model, q, dobs, dm; options=Options(), ms_func=multi_src_fg, kw...)
     # Number of sources and init result
     nsrc = try q.nsrc catch; dobs.nsrc end
     pool = _worker_pool()
@@ -111,7 +111,7 @@ function multi_src_fg!(G, model, q, dobs, dm; options=Options(), kw...)
     arg_func = i -> (model, q[i], dobs[i], dm, options[i])
     kw_func = i -> Dict(:illum=> illum, Dict(k => kw_i(v, i) for (k, v) in kw)...)
     # Distribute source
-    res = run_and_reduce(multi_src_fg, pool, nsrc, arg_func; kw=kw_func)
+    res = run_and_reduce(ms_func, pool, nsrc, arg_func; kw=kw_func)
     f, g = update_illum(res, model, :adjoint_born)
     f, g = as_vec(res, Val(options.return_array))
     G .+= g