Hi, all
I have a model which contains two parts. The first part “model1” takes one image and outputs a feature ‘model1_feat’. The second part “model2” takes the ‘model1_feat’ and another feature ‘input_feat’ as input, and generate the final output. I want to train this model on multi GPUs. I have writen the following code:
model1 = nn.DataParallel(model1).cuda()
model1_feat = model1(input_image)

model2 = nn.DataParallel(model2).cuda()
model2_feat = model2(model1_feat, input_feat)

But it does not work. The whole thread is blocked and the model can not generate any output. Can you help me.
BTW, the total model works fine on single card.
Thanks.