Multiple loss using DataParallel

Hi, I have multiple losses.
It worked well on single gpu.
However when I tried it using DataParallel, it stuck on loss.backward().

What should I consider when I use multiple loss on multiple gpus? Below code is a simple flow of my implementation.

criterion = nn.CrossEntropyLoss().cuda()
model = DataParallel(model).cuda()


model.train()
for i, (inputs, labels) in  enumerate(train_loader)

    labels = labels.cuda(non_blocking=True)    

    logits = model(inputs)

    loss_1 = criterion(logits[0], labels)             
    loss_2 = criterion(logits[1], labels)  

    loss = loss_1 + loss_2  
    loss.backward()


Thanks for your advice in advance!

Are you able to use nn.DataParallel with a single loss?

Yes, I think I managed to do it.

I found the problem I had.

os.environ['CUDA_LAUNCH_BLOCKING'] = '1'

This code was inserted to debug when I ran it on a single gpu. After I remove code above, my code works well.