Hi supposed I have two networks,
unet → for segmentation
resnet → for classification
How do I do joint training? is the code below correct?

def train_models(inputs, targets, labels):
    unet.train()
    resnet.train()
    segment_output = unet(inputs)
    output = resnet(segment_output)
   
    optimizer1.zero_grad()
    optimizer2.zero_grad()
    loss1 = MSE(segmentation_output, labels)
    loss2 = BCE(output, target)
    loss1.backward()
    loss2.backward()

    optimizer1.step()
    optimizer2.step()

Or this one is the correct way?
optimizer = optim.Adam([
    {'params': unet.parameters()},
    {'params': resnet.parameters(), 'lr': 1e-3}
], lr=1e-4)

def train_models(inputs, targets, labels):
    unet.train()
    resnet.train()
    segment_output = unet(inputs)
    output = resnet(segment_output)
   
    optimizer.zero_grad()
    loss1 = MSE(segmentation_output, labels)
    loss2 = BCE(output, target)
    loss = loss1 + loss2
    loss.backward()
    optimizer.step()

If I happened to have weighted loss, is the 2nd approach better?