I had a slightly different problem when training multiple models in multiple gpu in parallel. This is the only answer I found helpful so I hope I can bring it up here.
It can be shown by modifying the example based on the above suggestion. Here I want to run the processes (use_gpu) in all the devices in parallel.
import time
import torch
from torch.multiprocessing import Pool
torch.multiprocessing.set_start_method('spawn', force=True)


def use_gpu(ind):
    t = []
    for i in range(5):
        time.sleep(1)
        a = torch.randn(1000, 1000).cuda(ind)
        t.append(a)
    return t

if __name__ == "__main__":
    # torch.cuda.set_device(3)
    pool = Pool()
    result = []
    a = time.time()
    for i in range(torch.cuda.device_count()):
        result.append(pool.apply_async(use_gpu, (i,)))
    pool.close()
    pool.join()
    print("cost time :", time.time() - a)

However I couldn’t figure out why I got the error message
THCudaCheck FAIL file=/opt/conda/conda-bld/pytorch_1524590031827/work/aten/src/THC/THCTensorRandom.cu line=25 error=46 : all CUDA-capable devices are busy or unavailable
THCudaCheck FAIL file=/opt/conda/conda-bld/pytorch_1524590031827/work/aten/src/THC/THCTensorRandom.cu line=25 error=46 : all CUDA-capable devices are busy or unavailable
THCudaCheck FAIL file=/opt/conda/conda-bld/pytorch_1524590031827/work/aten/src/THC/THCTensorRandom.cu line=25 error=46 : all CUDA-capable devices are busy or unavailable

Thanks for the help.