[JacobiX]

https://pytorch.org/cppdocs/ This is a brief description of the API. It worked fine for us even for complex models.