Hi,
I would advice first to try and write your code without the for loop using pytorch ops. Even if you need extra computations/complex indexing, this is most likely going to be much faster.
If you forward of a single problem is very small, then the gpu is going to be very ineficient to run it, even if you parallelise the work of launching the jobs on the gpu.
You are not allowed to send tensors that require gradients to other process as it is very hard to track them and backpropagate properly through such move. You would need to detach your tensors and manually send data back and forth.