Hi,

I’m facing an issue which avoids me to train my model using sciml_train.

I have the following training data:

```
model = Normal(5., 0.1)
raw_data = rand(model, (1, 100))
x = rand(model, (1; 100))
u0 = [x; 0]
```

I have an ODE defined as:

```
function initial_value_problem(u::AbstractArray, p, t) #dynamics of z and logpz given in the paper
z = u[1:end-1]
f = re(p)
[f(z')'; -t_J(f, z)]
end
prob = ODEProblem(initial_value_problem, u0, tspan)
```

With f a simple neural network with weights p:

```
nn = Chain(Dense(1, 1, swish), Dense(1, 1))
p, re = Flux.destructure(nn)
```

I also define a function predict_adjoint as follows:

```
function predict_adjoint(p) #we want to solve ivp with adjoint method
_prob = remake(prob, u0 = u0, p=p)
return solve(_prob ,Tsit5(), #[u0,0f0]
saveat=0f0:0.1f0:10f0,sensealg=DiffEqFlux.InterpolatingAdjoint(
checkpointing=true)).u[end][1]
end
```

Then when I’m trying to compute gradient of a loss:

```
function loss_adjoint(p)
pz = Normal(0.0, 1.0)
preds = predict_adjoint(p)[end]
z0 = preds[1:end-1]
delta_logp = preds[end]
logpz = DistributionsAD.logpdf.(pz, z0)
logpx = logpz .- delta_logp
loss = -mean(logpx)
loss
end
```

I get the following error:

```
ERROR: MethodError: no method matching flatten(::Array{Float64,2})
```

Investigating into this error I found that it might come from step predict_adjoint in loss_adjoint.

I don’t understand why because loss_adjoint returns a scalar and every operation invoked inside should be differentiable according to Zygote so i should be able to compute gradient.

Do you know where I am wrong?

Thank you