import torch


def set_loss_weight(pivot_set,exp):
    """
    The weight of the k-th auxiliary loss: gamma_k = \max(0.01, (\frac{L_k}{L_K})^2)
    More details can be found in Section 3.2 in "The Shallow End: Empowering Shallower Deep-Convolutional Networks
    through Auxiliary Outputs": https://arxiv.org/abs/1611.01773.
    """


    base_weight = 0

    lr_weight = torch.zeros(len(pivot_set)).cuda()
    pivot_weight = lr_weight.clone()

    for i in range(len(pivot_set) - 1, -1, -1):
        temp_weight = max(pow(float(pivot_set[i]/pivot_set[-1]), exp), 0.01)

        base_weight += temp_weight
        pivot_weight[i] = temp_weight
        lr_weight[i] = base_weight

    return pivot_weight, lr_weight