wwiiiii:

Even though parallel_apply function runs the forward function of replica with torch.cuda.device(device of given inputs) context, I think it only affects where new tensors located, not the replica.


It’s not true. When we don’t specify devices when constructing nn.DataParallel instance (which is common case), devices are set as list of None 4, and then become list of current device 2, for all replicated modules.
But it doesn’t matter since each i-th (module, inputs) pair is already located on the GPU #i. So the result of calculation is also going to be located on i-th GPU, which makes with torch.cuda.device(device) meaningless. And that’s why we need the gather after the parallel_apply.
Please let me know if there’s any wrong point.