Sensitivity of an SDE with non-diagonal noise

Your noise rate prototype matrix isn’t converted. We should do that automatically in the library, but until we do this should work:

function loss_function(p)
    noise_prototype = zeros(eltype(p), 2, 2)
    sde_prob = SDEProblem(drift!, diffusion!, u0, tspan, p, noise_rate_prototype=noise_prototype)
    
    # Solve ensemble of SDEs
    ensemble_prob = EnsembleProblem(prob)
    ensemble_sol = solve(ensemble_prob, SRA1(), EnsembleThreads(), trajectories=num_trajectories, saveat=0.1)

    return sum(sum(sol[1, :] for sol in ensemble_sol.u))
end