diff --git a/benchmark/benchmarks.jl b/benchmark/benchmarks.jl
index 673dcba5..a751cdfc 100644
--- a/benchmark/benchmarks.jl
+++ b/benchmark/benchmarks.jl
@@ -52,6 +52,7 @@ icnf = ContinuousNormalizingFlows.construct(
         alg = OrdinaryDiffEqAdamsBashforthMoulton.VCABM(; thread = Static.True()),
         sensealg = SciMLSensitivity.InterpolatingAdjoint(;
             autodiff = true,
+            autojacvec = ZygoteVJP(),
             checkpointing = true,
         ),
     ),
@@ -124,6 +125,7 @@ icnf2 = ContinuousNormalizingFlows.construct(
         alg = OrdinaryDiffEqAdamsBashforthMoulton.VCABM(; thread = Static.True()),
         sensealg = SciMLSensitivity.InterpolatingAdjoint(;
             autodiff = true,
+            autojacvec = ZygoteVJP(),
             checkpointing = true,
         ),
     ),
diff --git a/examples/usage.jl b/examples/usage.jl
index c6838e0f..75588bd9 100644
--- a/examples/usage.jl
+++ b/examples/usage.jl
@@ -42,7 +42,11 @@ icnf = construct(
         abstol = eps(one(Float32)),
         maxiters = typemax(Int),
         alg = VCABM(; thread = True()),
-        sensealg = InterpolatingAdjoint(; autodiff = true, checkpointing = true),
+        sensealg = InterpolatingAdjoint(;
+            autodiff = true,
+            autojacvec = ZygoteVJP(),
+            checkpointing = true,
+        ),
     ), # pass to the solver
 )
 
diff --git a/src/utils.jl b/src/utils.jl
index 0527f19d..173c982a 100644
--- a/src/utils.jl
+++ b/src/utils.jl
@@ -9,7 +9,7 @@ function jacobian_batched(
     res = Zygote.Buffer(xs, size(xs, 1), size(xs, 1), size(xs, 2))
     for i in axes(xs, 1)
         ChainRulesCore.@ignore_derivatives z[i, :] .= one(T)
-        res[i, :, :] =
+        res[i, :, :] .=
             only(DifferentiationInterface.pullback(f, icnf.compute_mode.adback, xs, (z,)))
         ChainRulesCore.@ignore_derivatives z[i, :] .= zero(T)
     end
@@ -27,7 +27,7 @@ function jacobian_batched(
     res = Zygote.Buffer(xs, size(xs, 1), size(xs, 1), size(xs, 2))
     for i in axes(xs, 1)
         ChainRulesCore.@ignore_derivatives z[i, :] .= one(T)
-        res[:, i, :] = only(
+        res[:, i, :] .= only(
             DifferentiationInterface.pushforward(f, icnf.compute_mode.adback, xs, (z,)),
         )
         ChainRulesCore.@ignore_derivatives z[i, :] .= zero(T)
diff --git a/test/instability_tests.jl b/test/instability_tests.jl
index e8c52728..1b06e771 100644
--- a/test/instability_tests.jl
+++ b/test/instability_tests.jl
@@ -31,6 +31,7 @@ Test.@testset "Instability" begin
             alg = OrdinaryDiffEqAdamsBashforthMoulton.VCABM(; thread = Static.True()),
             sensealg = SciMLSensitivity.InterpolatingAdjoint(;
                 autodiff = true,
+                autojacvec = ZygoteVJP(),
                 checkpointing = true,
             ),
         ),
diff --git a/test/regression_tests.jl b/test/regression_tests.jl
index f0599b7f..22cc34ae 100644
--- a/test/regression_tests.jl
+++ b/test/regression_tests.jl
@@ -27,6 +27,7 @@ Test.@testset "Regression Tests" begin
             alg = OrdinaryDiffEqAdamsBashforthMoulton.VCABM(; thread = Static.True()),
             sensealg = SciMLSensitivity.InterpolatingAdjoint(;
                 autodiff = true,
+                autojacvec = ZygoteVJP(),
                 checkpointing = true,
             ),
         ),
diff --git a/test/smoke_tests.jl b/test/smoke_tests.jl
index 452cdc9d..3cba3715 100644
--- a/test/smoke_tests.jl
+++ b/test/smoke_tests.jl
@@ -130,6 +130,7 @@ Test.@testset "Smoke Tests" begin
                 alg = OrdinaryDiffEqAdamsBashforthMoulton.VCABM(; thread = Static.True()),
                 sensealg = SciMLSensitivity.InterpolatingAdjoint(;
                     autodiff = true,
+                    autojacvec = ZygoteVJP(),
                     checkpointing = true,
                 ),
             ),