diff --git a/Project.toml b/Project.toml
index 2ad8780ea..edd76b540 100644
--- a/Project.toml
+++ b/Project.toml
@@ -39,7 +39,7 @@ DiffResults = "1"
 DifferentiationInterface = "0.6, 0.7"
 Distributions = "0.25.111"
 DocStringExtensions = "0.8, 0.9"
-DynamicPPL = "0.39"
+DynamicPPL = "0.39, 0.40"
 Enzyme = "0.13"
 FillArrays = "1.3"
 Functors = "0.4, 0.5"
diff --git a/ext/AdvancedVIDynamicPPLExt.jl b/ext/AdvancedVIDynamicPPLExt.jl
index 89de7b291..eaff70ae7 100644
--- a/ext/AdvancedVIDynamicPPLExt.jl
+++ b/ext/AdvancedVIDynamicPPLExt.jl
@@ -48,7 +48,7 @@ end
 function logdensity_impl(
     params, model::DynamicPPL.Model, loglikeadj::Real, varinfo::DynamicPPL.AbstractVarInfo
 )
-    vi = DynamicPPL.unflatten(varinfo, params)
+    vi = DynamicPPL.unflatten!!(varinfo, params)
     _, vi = DynamicPPL.evaluate!!(model, vi)
     loglike = DynamicPPL.getloglikelihood(vi)
     logprior = DynamicPPL.getlogprior(vi)
diff --git a/test/Project.toml b/test/Project.toml
index a1c980829..285b91a9b 100644
--- a/test/Project.toml
+++ b/test/Project.toml
@@ -28,7 +28,7 @@ ADTypes = "0.2.1, 1"
 DiffResults = "1"
 DifferentiationInterface = "0.6, 0.7"
 Distributions = "0.25.111"
-DynamicPPL = "0.39"
+DynamicPPL = "0.39, 0.40"
 Enzyme = "0.13, 0.14, 0.15"
 FillArrays = "1.6.1"
 ForwardDiff = "0.10.36, 1"