Ah, we do not have a public API for exposing a static backward graph, as PyTorch relies on dynamic autograd for doing automatic differentiation. We do have an internal API for symbolic differentiation (see torch/csrc/jit/runtime/autodiff.cpp which you can play with, but it is not complete and we don’t have any stability guarantees about it