diff --git a/Project.toml b/Project.toml
index 550839a..eeacdd9 100644
--- a/Project.toml
+++ b/Project.toml
@@ -7,6 +7,7 @@ version = "0.1.0"
 AbbreviatedStackTraces = "ac637c84-cc71-43bf-9c33-c1b4316be3d4"
 BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
 CairoMakie = "13f3f980-e62b-5c42-98c6-ff1f3baf88f0"
+ChainRules = "082447d4-558c-5d27-93f4-14fc19e9eca2"
 Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
 Distributed = "8ba89e20-285c-5b6f-9357-94700520ee1b"
 Downloads = "f43a241f-c20a-4ad4-852c-f6b1247861c6"
@@ -38,7 +39,6 @@ SnoopPrecompile = "66db9d55-30c0-4569-8b51-7e840670fc0c"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 TimerOutputs = "a759f4b9-e2f1-59dc-863e-4aeb61b1ea8f"
 Tullio = "bc48ee85-29a4-5162-ae0b-a64e1601d4bc"
-Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [compat]
 AbbreviatedStackTraces = "0.1.7"
@@ -66,10 +66,10 @@ PyCall = "1.9"
 PyPlot = "2.11"
 Revise = "3.1"
 SciMLSensitivity = "7.20"
+ChainRules = "1.50"
 SnoopPrecompile = "1.0.3"
 TimerOutputs = "0.5.22"
 Tullio = "0.3"
-Zygote = "0.6"
 julia = "1.7"
 
 [extras]
diff --git a/data/gdir_refs.jld2 b/data/gdir_refs.jld2
deleted file mode 100644
index 11a88a9..0000000
Binary files a/data/gdir_refs.jld2 and /dev/null differ
diff --git a/data/gdir_refs_(2010.0, 2015.0).jld2 b/data/gdir_refs_(2010.0, 2015.0).jld2
new file mode 100644
index 0000000..d92e5cc
Binary files /dev/null and b/data/gdir_refs_(2010.0, 2015.0).jld2 differ
diff --git a/data/gdir_refs_(2010.0, 2015.0)_MB1.jld2 b/data/gdir_refs_(2010.0, 2015.0)_MB1.jld2
new file mode 100644
index 0000000..997135c
Binary files /dev/null and b/data/gdir_refs_(2010.0, 2015.0)_MB1.jld2 differ
diff --git a/data/gdir_refs_(2010.0, 2015.0)_MB2.jld2 b/data/gdir_refs_(2010.0, 2015.0)_MB2.jld2
new file mode 100644
index 0000000..406c133
Binary files /dev/null and b/data/gdir_refs_(2010.0, 2015.0)_MB2.jld2 differ
diff --git a/data/gdir_refs_(2010.0, 2015.0)_MB3.jld2 b/data/gdir_refs_(2010.0, 2015.0)_MB3.jld2
new file mode 100644
index 0000000..d92e5cc
Binary files /dev/null and b/data/gdir_refs_(2010.0, 2015.0)_MB3.jld2 differ
diff --git a/scripts/dhdt_plots.jl b/scripts/dhdt_plots.jl
index ed77e0d..873203b 100644
--- a/scripts/dhdt_plots.jl
+++ b/scripts/dhdt_plots.jl
@@ -8,15 +8,17 @@ import ODINN: fillZeros
 function make_plots()
 
     # plot_type = "only_H" # plot final H
-    # plot_type = "MB_diff" # differences between runs with different MB models
-    plot_type = "H_diff" # H - H₀
+    plot_type = "MB_diff" # differences between runs with different MB models
+    # plot_type = "H_diff" # H - H₀
     tspan = (2010.0, 2015.0) # period in years for simulation
 
     root_dir = dirname(Base.current_project())
 
     # Load forward simulations with different surface MB
     grefs = load(joinpath(root_dir, "data/gdir_refs_$tspan.jld2"))["gdir_refs"]
-    grefs_MBu1 = load(joinpath(root_dir, "data/gdir_refs_updatedMB1.jld2"))["gdir_refs"]
+    grefs_MBu1 = load(joinpath(root_dir, "data/gdir_refs_$(tspan)_MB1.jld2"))["gdir_refs"]
+    grefs_MBu2 = load(joinpath(root_dir, "data/gdir_refs_$(tspan)_MB2.jld2"))["gdir_refs"]
+    grefs_MBu3 = load(joinpath(root_dir, "data/gdir_refs_$(tspan)_MB3.jld2"))["gdir_refs"]
 
     n=4
     m=3
@@ -33,8 +35,7 @@ function make_plots()
         H = reverse(grefs[i]["H"]', dims=2)
         H₀ = reverse(grefs[i]["H₀"]', dims=2)
         H_MBu1 = reverse(grefs_MBu1[i]["H"]', dims=2)
-        # H = reverse(grefs[i]["H"])
-        # H_MBu1 = reverse(grefs_MBu1[i]["H"])
+        H_MBu2 = reverse(grefs_MBu3[i]["H"]', dims=2)
         if plot_type == "only_H"
             H_plot = H
             label = "Predicted H (m)"
@@ -42,7 +43,7 @@ function make_plots()
             H_plot = H .- H₀
             label = "H - H₀ (m)"
         elseif plot_type == "MB_diff"
-            H_plot = H .- H_MBu1
+            H_plot = H_MBu1 .- H_MBu2
             label="Surface mass balance difference (m)"
         end
         push!(MBdiffs, H_plot)
@@ -66,35 +67,6 @@ function make_plots()
 
     end # let
 
-    # hms = []
-    # for (gref, gref_MBu1) in zip(grefs, grefs_MBu1)
-    #     H = reverse(gref["H"], dims=1)
-    #     H_MBu1 = reverse(gref_MBu1["H"], dims=1)
-    #     # H = gref["H"]
-    #     # H_MBu1 = gref_MBu1["H"]
-    #     push!(hms, heatmap(H .- H_MBu1, 
-    #                         clims=(0.0,5.0),
-    #                         ylimits=(0, size(H)[1]),
-    #                         xlimits=(0, size(H)[2]),
-    #                         colorbar = false)
-    #         )
-    # end
-
-    # h2 = scatter([0,0], [0,1], clims=(0.0,5.0),
-    #                  xlims=(1,1.1), xshowaxis=false, yshowaxis=false, label="", colorbar_title="cbar", grid=false)
-
-
-    # l = @layout [grid(6,5) a{0.01w}]
-
-    # # Create the combined plot with the subplots and shared colormap
-    # p_dhdt = plot(hms..., h2,
-    #               size=(1800, 1200),
-    #               layout=l,
-    #               link=:all,
-    #               aspect_ratio=:equal)
-
-    # savefig(p_dhdt, joinpath(root_dir, "plots/MB/dhdt_MB_1"))
-
 end
 
 make_plots()
\ No newline at end of file
diff --git a/scripts/toy_model.jl b/scripts/toy_model.jl
index a57d29f..efb2677 100644
--- a/scripts/toy_model.jl
+++ b/scripts/toy_model.jl
@@ -24,7 +24,7 @@ ENV["GKSwstype"]="nul"
 
 function run_toy_model()
 
-    processes = 1
+    processes = 18
 
     # Enable multiprocessing
     ODINN.enable_multiprocessing(processes)
@@ -35,9 +35,9 @@ function run_toy_model()
     ODINN.set_run_spinup(false) # Run the spin-up random_MB = generate_random_MB(gdirs_climate, tspan; plot=false)n
     ODINN.set_use_spinup(false) # Use the updated spinup
     # Reference simulations
-    ODINN.set_create_ref_dataset(true) # Generate reference data for UDE training
+    ODINN.set_create_ref_dataset(false) # Generate reference data for UDE training
     # UDE training
-    ODINN.set_train(false)    # Train UDE
+    ODINN.set_train(true)    # Train UDE
     ODINN.set_retrain(false) # Re-use previous NN weights to continue training
     # Optimization method for differentiating the model
     ODINN.set_optimization_method("AD+AD")
@@ -69,7 +69,7 @@ function run_toy_model()
     if ODINN.create_ref_dataset[]
         println("Generating reference dataset for training...")
         # Compute reference dataset in parallel
-        mb_model = ODINN.TI_model_1(DDF=8.0/1000.0, acc_factor=1.0/1000.0) # in m.w.e.
+        mb_model = ODINN.TI_model_1(DDF=4.0/1000.0, acc_factor=2.0/1000.0) # in m.w.e.
         gdir_refs = @time "PDEs" generate_ref_dataset(gdirs, tspan, solver=RDPK3Sp35(), mb_model)
 
         println("Saving reference data")
@@ -89,7 +89,7 @@ function run_toy_model()
     if ODINN.train[]
         # Train iceflow UDE in parallel
         n_ADAM = 20
-        n_BFGS = 100
+        n_BFGS = 50
         batch_size = length(gdir_refs)
         # batch_size = 9
         UDE_settings = Dict("reltol"=>1e-7,
@@ -103,7 +103,7 @@ function run_toy_model()
             # optimizer = BFGS(initial_stepnorm=0.0001)
             optimizer = Adam(0.001)
             train_settings = (optimizer, n_ADAM, batch_size) # optimizer, epochs, batch_size
-            mb_model = ODINN.TI_model_1(DDF=6.0/1000.0, acc_factor=1.0/1000.0) # in m.w.e.
+            mb_model = ODINN.TI_model_1(DDF=8.0/1000.0, acc_factor=1.0/1000.0) # in m.w.e.
             iceflow_trained, UA_f, loss_history = @time "UDEs" train_iceflow_UDE(gdirs, gdir_refs,
                                                                         tspan, train_settings, θ_trained; 
                                                                         UDE_settings=UDE_settings,
diff --git a/src/ODINN.jl b/src/ODINN.jl
index 00ccf4f..465599e 100644
--- a/src/ODINN.jl
+++ b/src/ODINN.jl
@@ -12,7 +12,7 @@ using SciMLSensitivity
 using Optimization, Optim, OptimizationOptimJL
 import OptimizationOptimisers.Adam
 using IterTools: ncycle
-using Zygote: @ignore 
+using ChainRules: @ignore_derivatives
 using Base: @kwdef
 using Flux
 using Tullio
diff --git a/src/helpers/climate.jl b/src/helpers/climate.jl
index c116cc1..fdc481d 100644
--- a/src/helpers/climate.jl
+++ b/src/helpers/climate.jl
@@ -160,7 +160,7 @@ function trim_period(period, climate)
     return period
 end
 
-function partial_year(period::Type{<:Period}, float::AbstractFloat)
+function partial_year(period::Type{<:Period}, float)
     _year, Δ = divrem(float, 1)
     year_start = Date(_year)
     year = period((year_start + Year(1)) - year_start)
@@ -172,7 +172,7 @@ partial_year(float::AbstractFloat) = partial_year(Day, float)
 
 function get_longterm_temps(gdir::PyObject, tspan)
     climate = xr.open_dataset(joinpath(gdir.dir, "raw_climate_$tspan.nc")) # load only once at the beginning
-    dem = xr.open_rasterio(gdir.get_filepath("dem"))
+    dem = rioxarray.open_rasterio(gdir.get_filepath("dem"))
     apply_t_grad!(climate, dem)
     longterm_temps = climate.groupby("time.year").mean().temp.data
     return longterm_temps
diff --git a/src/helpers/iceflow.jl b/src/helpers/iceflow.jl
index 2fb44cd..bec1a67 100644
--- a/src/helpers/iceflow.jl
+++ b/src/helpers/iceflow.jl
@@ -33,7 +33,7 @@ function generate_ref_dataset(gdirs, tspan, mb_model; solver = RDPK3Sp35(), velo
     println("Running forward PDE ice flow model...\n")
     # Run batches in parallel
     A_noises = randn(rng_seed(), length(gdirs)) .* noise_A_magnitude
-    refs = @showprogress pmap((gdir, A_noise) -> batch_iceflow_PDE(gdir, A_noise, tspan, solver, mb_model; run_spinup=false, velocities=velocities), gdirs, A_noises)
+    refs = @showprogress map((gdir, A_noise) -> batch_iceflow_PDE(gdir, A_noise, tspan, solver, mb_model; run_spinup=false, velocities=velocities), gdirs, A_noises)
 
     # Gather information per gdir
     gdir_refs = get_gdir_refs(refs, gdirs)
@@ -359,10 +359,10 @@ function batch_iceflow_UDE(θ, UA_f, context, gdir, UDE_settings, mb_model)
         mb_model = TI_model_1(DDF=5.0/1000.0, acc_factor=1.2/1000.0)
     end
     # Initialize climate dataset
-    tstops, step = @ignore define_callback_steps(tspan)
+    tstops, step = @ignore_derivatives define_callback_steps(tspan)
     S_coords = context[13]
     dist_border = context[17]
-    climate = @ignore init_climate(gdir, tspan, step, context[9], S_coords)
+    climate = @ignore_derivatives init_climate(gdir, tspan, step, context[9], S_coords)
     # Callback  
     # Define stop times every one month
     stop_condition(u,t,integrator) = stop_condition_tstops(u,t,integrator, tstops) #closure
@@ -371,7 +371,7 @@ function batch_iceflow_UDE(θ, UA_f, context, gdir, UDE_settings, mb_model)
             S::Matrix{Float64} = context[9]
             S_coords = context[13]
             MB = context[15]
-            @ignore MB_timestep!(MB, mb_model, climate, S, S_coords, integrator.t, step)
+            @ignore_derivatives MB_timestep!(MB, mb_model, climate, S, S_coords, integrator.t, step)
             apply_MB_mask!(integrator.u, MB, MB_total, dist_border)
         end
         # Recompute A value
@@ -401,12 +401,12 @@ function batch_iceflow_UDE(θ, UA_f, context, gdir, UDE_settings, mb_model)
     testmode = context[16]
     V̄x_pred::Matrix{Float64}, V̄y_pred::Matrix{Float64} = avg_surface_V(context, H_pred, mean(climate.longterm_temps), "UDE", θ, UA_f; 
                                                                         testmode=testmode) # Average velocity with average temperature
-    rgi_id::String = @ignore gdir.rgi_id
+    rgi_id::String = @ignore_derivatives gdir.rgi_id
     H_V_pred = (H_pred, V̄x_pred, V̄y_pred, rgi_id)
 
     # println("Total MB for $rgi_id: ", sum(MB_total))
 
-    @ignore GC.gc() # Run the garbage collector to tame the RAM
+    @ignore_derivatives GC.gc() # Run the garbage collector to tame the RAM
 
     return H_V_pred
 end
@@ -418,11 +418,11 @@ function batch_iceflow_UDE_inplace(θ, UA_f, gdir, tspan; solver = RDPK3Sp35())
     mb_model = TI_model_1(DDF=5.0/1000.0, acc_factor=1.2/1000.0)
 
     # Initialize climate dataset
-    _, step = @ignore define_callback_steps(tspan)
+    _, step = @ignore_derivatives define_callback_steps(tspan)
     S = context[3]
     S_coords = context[32]
     dist_border = context[33]
-    climate = @ignore init_climate(gdir, tspan, step, S, S_coords)
+    climate = @ignore_derivatives init_climate(gdir, tspan, step, S, S_coords)
     # Define stop times every one month
     tstops, step = define_callback_steps(tspan)
     stop_condition(u,t,integrator) = stop_condition_tstops(u,t,integrator, tstops) #closure
@@ -432,7 +432,7 @@ function batch_iceflow_UDE_inplace(θ, UA_f, gdir, tspan; solver = RDPK3Sp35())
             MB = context[25]
             S = context[3]
             S_coords = context[32]
-            MB_timestep!(MB, mb_model, climate, S, S_coords, integrator.t, step)
+            @ignore_derivatives MB_timestep!(MB, mb_model, climate, S, S_coords, integrator.t, step)
             apply_MB_mask!(integrator.u, MB, MB_total, context)
         end
         # Recompute A value
diff --git a/src/helpers/initial_conditions.jl b/src/helpers/initial_conditions.jl
index b98ba64..e075563 100644
--- a/src/helpers/initial_conditions.jl
+++ b/src/helpers/initial_conditions.jl
@@ -137,7 +137,7 @@ function build_UDE_context_inv(gdir, gdir_ref, tspan; run_spinup=false)
     # Get evolved tickness and surface
     H = gdir_ref["H"]
 
-    @ignore begin
+    @ignore_derivatives begin
         glacier_gd = xr.open_dataset(gdir.get_filepath("gridded_data"))
         H₁ = glacier_gd.consensus_ice_thickness.data
         fillNaN!(H₀) # Fill NaNs with 0s to have real boundary conditions
diff --git a/src/helpers/inversions.jl b/src/helpers/inversions.jl
index 0335adb..dd2e0ab 100644
--- a/src/helpers/inversions.jl
+++ b/src/helpers/inversions.jl
@@ -145,7 +145,7 @@ function loss_iceflow_inversion(θ, UD, gdirs_climate, gdir_refs, context_batche
     # l_V = (l_Vx + l_Vy)/2
 
     # Plot V diffs to understand training
-    @ignore begin
+    @ignore_derivatives begin
         plot_V_diffs(gdirs_climate, gdir_refs, V_preds)
     end
 
diff --git a/src/helpers/oggm.jl b/src/helpers/oggm.jl
index 2e75fca..bd78be0 100644
--- a/src/helpers/oggm.jl
+++ b/src/helpers/oggm.jl
@@ -24,8 +24,10 @@ function oggm_config(working_dir=joinpath(homedir(), "Python/OGGM_data"); oggm_p
 
     # Multiprocessing 
     multiprocessing = $oggm_processes > 1 ? true : false
-    PARAMS["mp_processes"] = $oggm_processes
     PARAMS["use_multiprocessing"] = multiprocessing # Let's use multiprocessing for OGGM
+    if multiprocessing
+        PARAMS["mp_processes"] = $oggm_processes
+    end
 
     end # @eval ODINN
     end # @everywhere
diff --git a/src/helpers/utils.jl b/src/helpers/utils.jl
index caf5293..d92c0aa 100644
--- a/src/helpers/utils.jl
+++ b/src/helpers/utils.jl
@@ -230,7 +230,7 @@ function get_NN(θ_trained)
         Dense(10,3, x->softplus.(x)),
         Dense(3,1, sigmoid_A)
     )
-    Flux.f64(UA)
+    UA = Flux.f64(UA)
     # See if parameters need to be retrained or not
     θ, UA_f = Flux.destructure(UA)
     if !isempty(θ_trained)
diff --git a/test/test.jl b/test/test.jl
index 926e692..52bdd7f 100644
--- a/test/test.jl
+++ b/test/test.jl
@@ -14,7 +14,7 @@ using Test
 # using HDF5  
 # using JLD2
 # using OrdinaryDiffEq, DiffEqFlux
-# using Zygote: @ignore
+# using Zygote: @ignore_derivatives
 # using Flux
 # using Tullio, RecursiveArrayTools
 # using Infiltrator