diff --git a/src/training_strategies.jl b/src/training_strategies.jl
index df185e5a6..2c30e91c9 100644
--- a/src/training_strategies.jl
+++ b/src/training_strategies.jl
@@ -320,7 +320,7 @@ function get_loss_function(loss_function, lb, ub, eltypeθ, strategy::Quadrature
             # mean(abs2,loss_(x,θ), dims=2)
             # size_x = fill(size(x)[2],(1,1))
             x = adapt(parameterless_type(ComponentArrays.getdata(θ)), x)
-            sum(abs2, vec(loss_(x, θ)), dims = 2) #./ size_x
+            vec(sum(abs2, loss_(x, θ), dims = 2)) #./ size_x
         end
         integral_function = BatchIntegralFunction(integrand, max_batch = strategy.batch)
         prob = IntegralProblem(integral_function, lb, ub, θ)