How to create a multilayer perceptron using Flux

General Usage
,
1

Hi
I’m new with Julia and I have some problems with the creation of a multilayer perceptron.
I’ve tried following some guides about this topic but I couldn’t figure it out the solution.
My question is: How can i add an ODEs as just a layer in my multilayer perceptron like in the guided example?
Here my example of solving ODEs as just a single layer:

#DATASET GENERATION

using DifferentialEquations
using Plots
using Flux, DiffEqFlux
using Optim
l = 1.0
m = 1.0
g = 9.81

#input
input(t) = 0.1sin(t) #Nm

#parameters
p = [-3*g/(2l),3/(m*l^2)] #[-14.715,3.0]


#Original function
function pendolo(du,u,p,t)
   theta, omega = u
   du[1] = dtheta = omega
   du[2] = domega = p[1]*sin(theta) + p[2]*input(t)
end

#Initial conditions
theta_0 = 0.01
omega_0 = 0.0
u0 = [theta_0,omega_0]

tspan = (0.0,10.0)
t = 0.1 

#ODE
prob = ODEProblem(pendolo,u0,tspan,p) 
sol = solve(prob,alg_hits = [:stiff], saveat = t);
ode_data = Array(sol[1:2,:]);

# SINGLE LAYER DECLARATION

p_random = [-12.5,4.5] #Random parameters to fit
prob_to_fit = ODEProblem(pendolo,u0,tspan,p_random) 
pred() = solve(prob_to_fit,Tsit5(), saveat = t);

#TRAINIG

#Neural parameter
param = Flux.params(p_random)       
loss() = Flux.mse(pred(),ode_data)  
data = Iterators.repeated((), 1000)
opt = RADAM(0.1)

cb = function ()                     #callback function to observe training
 display(loss())
 
   pl = plot(sol,layout = (2,1),label = [" θ " " ω"])
   Plots.scatter!(solve(remake(prob_to_fit, p = p_random),Tsit5(),saveat = t),label = [" nn_ode θ " " nn_ode ω"])
 display(plot(pl))
end

# Display the ODE with the initial parameter values.
cb()
@time begin
Flux.train!(loss, param, data, opt, cb = cb)
end

The problems come when I want to add other layers to my neural structure, here I have the same dataset generation previously described with this:

p_random = Float64[-12.0,5.0]
diff = ODEProblem(pendolo,u0,tspan,p_random)
layers = Chain(p -> solve(diff,Tsit5(), saveat = t_step,p = p)[1:2,:],
                Dense(301,50,relu),
                Dense(50,2))

n_ode = NeuralODE(layers,tspan,Tsit5(),saveat = t_range,reltol=1e-9,abstol=1e-9)
ps = Flux.params(p_random)
pred = n_ode(u0)

I’ve already tried with different possible solution changing the dimensions of the layers but without a solution.
The purpose of this script is to estimate the parameters of a system using a grey-box approach [knowing the mathematical structure of the system and a partial indication of the parameters].

Sorry for my stupid question but this is my first time using Julia.
Thank you in advanced for the support.