Haibin Ren

Hi, I just simply count the runtime of the training steps in `main()`, using the command below ``` python3 total_time = 0 for step, data_i in zip( range(steps), dataloader((ts, labels)...

Yes, I add this command below the training steps like this. ```python3 for step, data_i in zip( range(steps), dataloader((ts, labels) + coeffs, batch_size, key=loader_key) ): bxe, acc, model, opt_state =...

Sorry for the confusion! Here is the code click here ```python3 import math import diffrax import equinox as eqx # https://github.com/patrick-kidger/equinox import jax import jax.nn as jnn import jax.numpy as...

Hi @johannahaffner, with newest packages, I got the same result with your code. I guess the different benchmark results between `make_step()` and `NeuralCDE.__call__()` are because of two reasons. 1. The...