Since you are leaving PyTorch, you would have to write a custom autograd.Function, as described here 24, and also implement the backward pass manually.
Autograd won’t be able to track the operations in the other process.