XlaRuntimeError: INTERNAL: Generated function failed: CpuCallback error: RuntimeError: The maximum number of solver steps was reached. Try increasing `max_steps`.

Hi Patrick!

I'm training a simple Neural CDE model for use in time series regression, and I keep coming across the same error again and again. I know this error (RuntimeError: The maximum number of solver steps was reached. Try increasing max_steps) has come up in several issues, but I was hoping I could get some help with resolving this as I was unable to pinpoint why it's arising. I'm making use of the exact same Func and NeuralCDE classes as provided in the Neural CDE example in the Diffrax documentation; only things I changed is making use of the following loss and make_step functions:

#@eqx.filter_jit
def loss(model, ti, label_i, coeff_i):
    pred = jax.vmap(model)(ti, coeff_i)
    # huber loss
    errors = pred - label_i
    abs_errors = jnp.abs(errors)
    quadratic = jnp.minimum(abs_errors, 1)
    linear = abs_errors - quadratic
    rmse = jnp.sqrt(jnp.mean(errors ** 2))
    return jnp.mean(0.5 * quadratic ** 2 + linear), rmse, abs_errors

grad_loss = eqx.filter_value_and_grad(loss, has_aux=True)

#@eqx.filter_jit
def make_step(model, batch, opt_state):
    # ts, coeff_i, ys = jnp.array(batch['ts']), jnp.array(batch['coeffs']), jnp.array(batch['ys'])
    ts = jnp.array(batch['ts'], dtype=jnp.float64)
    coeff_i = jnp.array(batch['coeffs'], dtype=jnp.float64)
    ys = jnp.array(batch['ys'], dtype=jnp.float64)
    (huber, rmse, mae), grads = grad_loss(model, ts, ys, coeff_i)
    updates, opt_state = optim.update(grads, opt_state)
    model = eqx.apply_updates(model, updates)
    return huber, rmse, mae, model, opt_state

and making use of a PyTorch dataset class with custom dumpy collate function to handle dictionaries and the NumPy data. Any idea of where this error might be coming from?

Here is a link to my full code in case this might not be sufficient to help figure out what's wrong: https://pastebin.com/L4BM0yT8

Any help with this would be greatly appreciated!

Start with the Neural CDE example in the Diffrax documentation and compare it with the reported loss and make_step functions. Review the full code at the linked Pastebin, including the PyTorch dataset and custom collate function, and trace where the solver reaches max_steps. Done means identifying the cause of the CpuCallback error and confirming a working training run.