Lowest Possible Memory For Neural ODE

Specific Domains Machine Learning
,
1

For a neural ode, what are the best kwargs to pass to the solve function that make it use the lowest possible memory in both direct and AD usage?
I use:

kwargs = Dict(
    :alg_hints => [:nonstiff, :memorybound],
    :dense => false,
    :save_everystep => false,
    :save_on => false,
    :calck => false,
    :alias_u0 => true,
    :verbose => false,
    :merge_callbacks => true,
    :wrap => Val(true),
    :alg => BS3(),
    :sensealg => BacksolveAdjoint(;
        autodiff = true,
        autojacvec = ZygoteVJP(),
        checkpointing = true,
    ),
    :reltol => 1.0f-2 + eps(1.0f-2),
)

Can it be any better?

2

For the alg you can use a low memory RK method.