Hi, did you find a solution here, I find the similar question https://discuss.pytorch.org/t/how-to-parallelize-a-loop-over-the-samples-of-a-batch/32698 and https://discuss.pytorch.org/t/running-multiple-modules-in-parallel/58164/2, both answers give a way of reformulate the problem, however they did not solve the problem of multi-processing inside one GPU inside forward function.
2 Likes