See https://github.com/juntang-zhuang/TorchDiffEqPack/blob/master/test_code/three_body_problem_dense.py