losses.py

def create_loss(losses, args):

    alpha = args.SOLVER.ALPHA
    loss_backbone = losses[0]
    loss_head = losses[1]

    def loss(embeddings, est_difficulties,labels, difficulties, score_keeper):
        labels = labels.reshape(-1)
        if score_keeper.is_training == 'backbone':
            return loss_backbone(embeddings, labels)
        if score_keeper.is_training == 'head':
            return loss_head(est_difficulties, difficulties)
        if score_keeper.is_training == 'combined':
            return alpha * loss_backbone(embeddings, labels) + (1-alpha)*loss_head(est_difficulties, difficulties)
    return loss