Hmmm the think is that batches are distributed properly at the time of calling forward.
When you use dataparallel a copy of the model is created on each gpu and batch is distributed among them.
For example i see that inside montecarlo you have a reward variable initialized as zeros, but that reward is not properly allocated on its corresponding gpu.
Think that model parallel only allocates those variables which are defined before calling dataparallel. If you define a variable during the forward pass, you are the one in charge of allocating it in the proper gpu.
In the context of dataparallel, you have to softcode those variables allocating them into a device whose id depends on a variable or module device.