title: PyTorch allocates more memory on the first available GPU (cuda:0)
body: <p>As a part of the reinforcement learning training system, I am training four policies in parallel using four GPUs. For each model, there are two processes - the actor and the learner, which only use their specific GPU (e.g. actor and learner corresponding to model #2 only use GPU #2 for all their tensors). Actor and learner share the model layers via torch's <code>share_memory_()</code>.
  Since the four training "subsystems" are completely symmetric, I would expect them to use the exact same amount of GPU memory on each of the four GPUs. In practice, however, I see a lot more GPU memory allocated on the first GPU (<code>cuda:0</code>).</p>
  
  <p>It feels like all the memory sharing is somehow done via GPU #0. Is there a way to fix this?</p>
  
  <p>So far, I tried setting <code>CUDA_VISIBLE_DEVICES</code> in child processes by explicitly altering <code>os.environ</code> in the process <code>start</code> function. This does not seem to have any effect, probably because child processes are forked from the main process, where PyTorch CUDA is already initialized, and envvars just seem to be ignored at this point.</p>
  
  <p><a href="https://i.stack.imgur.com/HCzXY.png" rel="nofollow noreferrer"><img src="https://i.stack.imgur.com/HCzXY.png" alt="enter image description here"></a></p>