Sure!
The following is the training code where I get the outputs, compute losses and call.backward()
train_data = utils.IntensityLoader(data_path=TRAIN_DATA_PATH, batch_size=BATCH_SIZE,
                                                       index=i).loaded_data()
                    inputs, truth = train_data[0][0], train_data[0][1]
                    inputs = inputs.to(device)
                    truth = truth
                    truth = truth.float()
                    optimizer.zero_grad()
                    outputs1, outputs2, outputs3, outputs4, outputs5 = au_net(inputs)
                    
                    # I'm putting the outputs on CPU to save GPU memory
                    outputs1 = outputs1.cpu()
                    outputs2 = outputs2.cpu()
                    outputs3 = outputs3.cpu()
                    outputs4 = outputs4.cpu()
                    outputs5 = outputs5.cpu()

                    corr_factor = corr_net(inputs).reshape([BATCH_SIZE, 2, 5]).cpu()
                    corr_factor = torch.bmm(corr_factor, corr_factor.permute(0, 2, 1))
                    del_corr = torch.bmm(del_cor, del_cor.permute(0, 2, 1))
                    """
                    loss_2 is the CorrNet loss - this is a separate branch network with custom loss function
                    """
                    loss_2_1 = torch.Tensor(2, 2).fill_(-1.0) * (corr_factor + del_corr) * torch.bmm(
                        (outputs1 - train_AU_mean.float()).unsqueeze(1),
                        (outputs1 - train_AU_mean.float()).unsqueeze(1).permute(0, 2, 1))

                    loss_2_2 = torch.Tensor(2, 2).fill_(-1.0) * (corr_factor + del_corr) * torch.bmm(
                        (outputs2 - train_AU_mean.float()).unsqueeze(1),
                        (outputs2 - train_AU_mean.float()).unsqueeze(1).permute(0, 2, 1))

                    loss_2_3 = torch.Tensor(2, 2).fill_(-1.0) * (corr_factor + del_corr) * torch.bmm(
                        (outputs3 - train_AU_mean.float()).unsqueeze(1),
                        (outputs3 - train_AU_mean.float()).unsqueeze(1).permute(0, 2, 1))

                    loss_2_4 = torch.Tensor(2, 2).fill_(-1.0) * (corr_factor + del_corr) * torch.bmm(
                        (outputs4 - train_AU_mean.float()).unsqueeze(1),
                        (outputs4 - train_AU_mean.float()).unsqueeze(1).permute(0, 2, 1))

                    loss_2_5 = torch.Tensor(2, 2).fill_(-1.0) * (corr_factor + del_corr) * torch.bmm(
                        (outputs5 - train_AU_mean.float()).unsqueeze(1),
                        (outputs5 - train_AU_mean.float()).unsqueeze(1).permute(0, 2, 1))

                    loss_2 = (loss_2_1 + loss_2_2 + loss_2_3 + loss_2_4 + loss_2_5) / 5.0

                    corr_loss = torch.mean(loss_2[torch.triu(torch.ones(BATCH_SIZE, 2, 2)) == 1])

                    print ("corr loss", corr_loss)
                    loss_list = list()
                    loss = torch.zeros(1)
                    
                    # calculating losses for every head
                    for index in range(len(truth)):
                        sub_loss_1 = AU_criterion(outputs1[index].unsqueeze(0), truth[index][0].unsqueeze(0))
                        sub_loss_2 = AU_criterion(outputs2[index].unsqueeze(0), truth[index][1].unsqueeze(0))
                        sub_loss_3 = AU_criterion(outputs3[index].unsqueeze(0), truth[index][2].unsqueeze(0))
                        sub_loss_4 = AU_criterion(outputs4[index].unsqueeze(0), truth[index][3].unsqueeze(0))
                        sub_loss_5 = AU_criterion(outputs5[index].unsqueeze(0), truth[index][4].unsqueeze(0))
                        loss += sub_loss_1 + sub_loss_2 + sub_loss_3 + sub_loss_4 + sub_loss_5

                    loss = loss + 4.0 + loss_weight * corr_loss
                    # loss = loss * 10e-9  # attempted to scale the loss to a smaller value. That did not work
                    loss.backward()
                    print ("total loss", loss)
                    optimizer.step()
                    running_loss += loss.item()

The following is how I define the optimizer:
params = list(au_net.parameters()) + list(corr_net.parameters())
    optimizer = optim.Adam(params, lr=LEARNING_RATE, betas=(0.9, 0.999), eps=0.1)

au_net is a VGG setup with 5 FC heads for 5 outputs. Each head maps 4096x1. I later reduced that to 24x1 to reduce the size of the network. But that did not work either
corr_net is a CNN setup with a FC layer at the end
I am using MSE to compute the loss for the au_net. That is defined as the AU_criterion in the code