Error in parameter estimation

I have been trying out some simple parameter estimation problem for ODEs. In my code, I specify save_idxs = 2 in calculating the loss, I get this error: MethodError: no method matching vec(::Float32 . I need your help. Thanks.
Here is my full code.
`
using Flux, DiffEqFlux, DifferentialEquations, Plots
function f(du,u,p,t)
S,I,R = u
du[1] = -p[1]*u[1]*u[2]
du[2] = p[1]*u[1]*u[2] - p[2]*u[2]
du[3] = p[2]*u[2]
end

initialise and have arange between 0 and 1.5 for time-t

tbegin=0.0
tend=15
tstep=1
trange = tbegin:tstep:tend
u0 =Float32[738.0,1.0,0.0]
tspan = (tbegin,tend)

t_f = collect(3:14)
I_data=Float32[25.0,75.0,227.0,296.0,258.0,236.0,192.0,126.0,71.0,28.0,11.0,7.0];

p = [0.002,0.3] # set initial values of parameters
prob = ODEProblem(f, u0, tspan, p)

sol=solve(prob,Tsit5(),reltol = 1e-12 , saveat=t_f,save_idxs = 2)
plot(sol)

function loss(p)
sol = solve(prob, Tsit5(),reltol = 1e-12 , saveat=t_f, save_idxs = 2)
loss = sum(abs2, I_data .- sol)
return loss, sol
end

callback = function (p, l, pred)
display(l)

Tell sciml_train to not halt the optimization. If return true, then

optimization stops.

return false
end

using Optimization
adtype = Optimization.AutoZygote()
optf = Optimization.OptimizationFunction((x, p) → loss(x), adtype)
optprob = Optimization.OptimizationProblem(optf, p)
result_ode = Optimization.solve(optprob,
ADAM(0.1),
callback = callback,
maxiters = 300)
`

Duplicate of Error in parameter estimation of ODEs