adapted from `https://github.com/rtqichen/torchdiffeq/tree/master/torchdiffeq`