So currently this isn’t supported.



 jonathanking:

RuntimeError: Cowardly refusing to serialize non-leaf tensor which requires_grad, since autograd does not support crossing process boundaries. If you just want to transfer the data, call detach() on the tensor before serializing (e.g., putting it on the queue).


You could implement forward and backward separately and stick it in a autograd.Function. Then inside the forward, you don’t have grad-requiring things.
For a loss function in particular, you could compute the backward inside the forward as well (you can turn also autograd back on inside the forward with torch.enable_grad).
Then in the backward you just hand back the pre-computed result. (This is what CuDNN’s CTC Loss implementation does.)
It might be more efficient in terms of development time as well as runtime to drop into C++ and use PyTorch’s parallel_for, though.
Note that PyTorch tries to use intra-op parallelism, you would not want to “overbook” your cores.
Best regards
Thomas