def main():
    best_prec1 = 0
    lr = 0.0001
    epochs = 20
 
    model = vgg.vgg16(num_classes=1000, pretrained=True)

    model.cuda()

    criterion_A = nn.CrossEntropyLoss().cuda()    # loss of taskA
    criterion_B = nn.CrossEntropyLoss().cuda()    # loss of taskB
    
    #  learnable loss weight param
    log_sigma_A = torch.nn.Parameter(torch.Tensor([1]))     
    log_sigma_B = torch.nn.Parameter(torch.Tensor([1]))

    weight_list = []
    bias_list = []
    last_weight_list = []
    last_bias_list = []
    loss_weight_list = [log_sigma_A, log_sigma_B]

    for name, value in model.named_parameters():
        if 'classifier' in name:
            #print(name)
            if 'weight' in name:
                last_weight_list.append(value)
            elif 'bias' in name:
                last_bias_list.append(value)
        else:
            if 'weight' in name:
                weight_list.append(value)
            elif 'bias' in name:
                bias_list.append(value)

    optimizer = torch.optim.SGD([{'params': weight_list, 'lr': lr},
                     {'params': bias_list, 'lr': lr * 2},
                     {'params': last_weight_list, 'lr': lr * 10},
                     {'params': last_bias_list, 'lr': lr * 20},
                     {'params': loss_weight_list, 'lr': lr}], momentum=0.9, weight_decay=0.0005, nesterov=True)

    for epoch in range(0, epochs):
        scheduler.step(epoch)
        # train for one epoch
        train(train_loader, model, criterion_A, criterion_B, optimizer, epoch, log_sigma_A, log_sigma_B)

def train(train_loader, model, criterion_A, criterion_B, optimizer, epoch, log_sigma_A, log_sigma_B):
            ....
        log_sigma_A = log_sigma_A.cuda()
        log_sigma_B = log_sigma_B.cuda()
        sigma_A = torch.Tensor.exp(log_sigma_A)
        sigma_B = torch.Tensor.exp(log_sigma_B)

        loss = (1/(2*sigma_A))*loss_A+ (1/(2*sigma_B)) * loss_B + log_sigma_A + log_sigma_B
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

log_sigma_A.grad is None     log_sigma_B.grad is None