Woops

JuliaDiff · Oct 10, 2024 · 720be01 · 720be01
1 parent 85d3195
commit 720be01
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/DifferentiationInterface/src/second_order/hessian.jl b/DifferentiationInterface/src/second_order/hessian.jl
@@ -111,14 +111,16 @@ function hessian(
         f, hvp_prep, backend, x, batched_seeds[1], contexts...
     )
 
-    hess = mapreduce(hcat, eachindex(batched_seeds)) do a
+    hess_blocks = map(eachindex(batched_seeds)) do a
         dg_batch = hvp(f, hvp_prep_same, backend, x, batched_seeds[a], contexts...)
         block = stack_vec_col(dg_batch)
         if N % B != 0 && a == lastindex(batched_seeds)
             block = block[:, 1:(N - (a - 1) * B)]
         end
         block
     end
+
+    hess = reduce(hcat, hess_blocks)
     return hess
 end