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