Hi guys, currently I’m trying to implement a CNN which has two outputs(like the figure below).

未命名.png315×514 8.94 KB

What I want is that Label1 will optimize CNN1 and Label2 will optimize CNN2. As for Shared CNN, it will be optimized by both Label1 and Label2. I’m wondering if it will work if the code is as follows:
output1, output2 = CNN(img)

loss1 = F.kl_div(output1, label1)
loss2 = F.kl_dic(output2, label2)

total_loss = loss1+loss2

optimizer.zero_grad()    # optimizer has declared before
total_loss.backward(retain_graph=True)
optimizer.step()

If it will not behave what I desire, can someone teach me or give me some hint of how to achieve the function? Thanks very much!