How to do multi-task training?

vision
7

I use nn.CrossEntropyLoss() as my criterion function, and the label as follow:

image

And code about loss is as follow:

# forward
outputs = model(inputs)
loss = 0

for lx in range(len(outputs)):
    tmp_loss = criterion(outputs[lx], labels[:, lx])
    loss += tmp_loss

# backward + optimize only if in training phase
if phase == 'train':
    loss.backward()
    optimizer.step()

The main problem here is how to add the losses of labels I have, because 0 stands for No Label and Invisible, it is not easy to differ.

Should I use other criterion function?