# Failure to calculate the gradient of ODE with the RHS as an argument: "LoadError: TypeError: in TrackedReal, in V, expected V<:Real, got Type{Any}"

Dear Forum,

I am trying to add ODE into machine learning pipeline and calculate the gradient of an ode layer that takes the RHS as an argument. We tried to feed the RHS, in terms of inputFunction(t), into ODE (i.e. fiip function). In fiip, inputFunction(t) is extracted from p, which is fed in function model(inputdata).
The code is as below.

``````using Flux
using DifferentialEquations, DiffEqSensitivity

parameter = [0.5]

function fiip(du, u, p, t)
p1, inputFunction = p
du[1] = dx = p1 * u[1] + inputFunction(t)
du[2] = dy = -p1 * u[2]
end

function model(inputdata)
inputFunction(t) = inputdata * t

p = [parameter[1], inputFunction]

u0 = [1.0; 1.0];
tspan = (0.0, 10.0)

concrete_solve(
ODEProblem(fiip, u0, tspan, p),
Tsit5(),
u0,
p,
saveat = 0:0.1:10,
)[1,:]
end

loss(x, y) = Flux.mse(model(x), y) # Loss function

x = 10
y = 100 * sin.(0:0.1:10)

gs = Flux.gradient(() -> loss(x, y), params(parameter))

``````

When the code executes to the last line, i.e. calculate the gradient of ODE with respect to the designated parameter, such error pops up:

LoadError: TypeError: in TrackedReal, in V, expected V<:Real, got Type{Any}

We’d be highly grateful for any insight into the possible origin and solution to this problem.
Also, will it help if we use callbacks on the ODE solver instead?

Are you using Julia v1.4?

Version 1.4.1

Oh, the issue is

is not an array of numbers. The adjoint methods need that the parameter vector is an array of values you can compute on. You can just use a closure for the anonymous function