import torch


def weighted_cross_entropy(input, target, weights):
    """ Compute cross entropy weighted per example
    got most ideas from https://github.com/pytorch/pytorch/issues/563

    Args:
        input (torch tensor): (N,C) log probabilities e.g. F.log_softmax(y_hat)
        target (torch tensor): (N) where each value is 0 <= targes[i] <=C-1
        weights (torch tensor): (N) per-example weight

    Return:
        nlhd: negative loglijelihood
    """
    # compute (N,) tensor of the log probabilities of the correct classes
    logpy = torch.gather(input, 1, target.view(-1,1)).view(-1)

    # compute the negative log-likelihood loss:
    nlhd = -(logpy * weights).mean()
    # nlhd = -logpy.mean()

    return nlhd