Training multiple pytorch models on GPUs
I'm trying to implement something with pytorch.
I have 2 GPUs and I want to train 2 models as below:
model0 = Mymodel().to('cuda:0')
model1 = Mymodel().to('cuda:1')
opt0 = torch.optim.Adam(model0.parameters(), lr=0.01)
opt1 = torch.optim.Adam(model0.parameters(), lr=0.01)

# 1.Forward data into model0 on GPU0
out = model0(x.to('cuda:0'))

# 2.Calculate the loss on model0, update model0's parameters
model0.loss.backward()
opt0.step()
opt0.zero_grad()

# 3.Use model0's output as input of model1 on GPU1
out = model1(out.to('cuda:1'))

# 4.Calculate the loss on model1, update model1's parameters
model1.loss.backward()
opt1.step()
opt1.zero_grad()

I want to train them simultaneously to speed up the whole procedure, but I think the code now will wait step 2(or 4) finished and finally do step 3(or 1). How can I implement my idea? Or which technique is I need(e.g. model parel, thread, multiprocessing...)?
I've consider some article like this, but I think there is some worng with the result, and I think it actually doesn't train models simultaneously.

You have a strong dependency between the 2 models, the 2nd one always needs the output from the previous one, so that part of the code will always be sequential.
I think you might need some sort of multiprocessing (take a look at torch.multiprocessing) or some kind of queue, where you can store the output from the first model.
