I have a CNN network that consists of multiple small Sequential modules and each of the cnn modules have different output neurons, which i return in the end. Something like this:
class CNN:
def __init__(self):
self.cnn_module = nn.Sequential() # contains CNN layers
self.network1 = nn.Sequential() # Dense layers
self.network2 = nn.Sequential() # Dense layers
self.network3 = nn.Sequential() # Dense layers

def forward(self, x):
x = self.cnn_module(x)
x = x.flatten(batch_size, -1)
out_1 = self.network1(x)
out_2 = self.network2(x)
out_3 = self.network3(x)
return [out_1, out_2, out_3]

I was confused on how to calculated loss of each of the networks and then backpropogate the losses together. Would something like this work:
criterion = nn.CrossEntropyLoss()
model = CNN()
output = model(images)
loss_1 = criterion(output[0], label_0)
loss_2 = criterion(output[1], labels_1)
loss_3 = criterion(output[2], labels_2)
loss = loss_1+loss_2+loss_3
loss.backward()
optimizer.step()