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]

function model(inputdata)
    inputFunction(t) = inputdata * t
    p = [parameter[1], inputFunction]
    u0 = [1.0; 1.0];
    tspan = (0.0, 10.0)
        ODEProblem(fiip, u0, tspan, p),
        saveat = 0:0.1:10,
        sensealg = QuadratureAdjoint(),

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

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

# Gradient calculation
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

Thanks so much for your informative comments. It worked for us!