Multi-Task Model


I have a problem on how to train a multi-task model. Suppose I have three classification tasks that need to be leaned simultaneously and I tend to use torch.nn.CrossEntropyLoss. Do I need to create criterion for each task, like:

criterion1 = nn.CrossEntropyLoss().cuda()
criterion2 = nn.CrossEntropyLoss().cuda()
criterion3 = nn.CrossEntropyLoss().cuda()

loss1 = criterion1(out1, target1)
loss2 = criterion2(out2, target2)
loss3 = criterion3(out3, target3)

or just create one criterion and it can be used for each task, like:

criterion = nn.CrossEntropyLoss().cuda()

loss1 = criterion(out1, target1)
loss2 = criterion(out2, target2)
loss3 = criterion(out3, target3)