The solution I would suggest is exactly what @Maxime_Louis said - when returning tensors from subgraphs, return the gradients too. There’s no way to make autograd work across processes without terrible hacks and a performance penalty - they don’t share the same address space.
If your network is so tiny I’d just go with using the CPU. Also, I would try to write batched code for the network so you don’t need to paralellize over elements, but take the input as a single large batch of independent values and return a batch of results. If you don’t see how to do that, try going with threads, because that will be visible to autograd.