diff --git a/Project.toml b/Project.toml
index 5e877a7..b7765fa 100644
--- a/Project.toml
+++ b/Project.toml
@@ -18,7 +18,7 @@ JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
 [compat]
 julia = "1.9"
 FFTW = "1"
-JLD2 = "0.4"
+JLD2 = "0.5"
 KernelAbstractions = "0.9.2"
 TimerOutputs = "0.5"
 ReadVTK = "0.2"
diff --git a/src/driver.jl b/src/driver.jl
index b956fe4..001e030 100644
--- a/src/driver.jl
+++ b/src/driver.jl
@@ -28,6 +28,7 @@ mutable struct LLG_CPP{T<:AbstractFloat} <: Driver
     integrator::Integrator
     tol::Float64
     p::Tuple{Real,Real,Real}
+    ufun::Function
 end
 
 mutable struct LLG_STT{T<:AbstractFloat} <: Driver
@@ -127,7 +128,7 @@ function create_driver(driver::String, integrator::String, n_total::Int64)
         if integrator == "DormandPrinceCayley"
             error(@sprintf("Unsupported combination driver %s and %s.", driver, integrator))
         end
-        return LLG_CPP(T(0.5), T(0), T(2.21e5), aj, T(0), dopri5, tol, p)
+        return LLG_CPP(T(0.5), T(0), T(2.21e5), aj, T(0), dopri5, tol, p, fun)
     end
     return nothing
 end