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