import torch


def cross_entropy(y_hat, y):
    """
    Compute the cross entropy as the loss of some nets

    Parameters
    ----------
    y_hat : [tensor]
        the prediction value
    y : [tensor]
        the real labels

    Returns
    -------
    [tensor]
        the vector of every cross entropy loss. They are negetive to be optimized
    """
    # here we use gather() to get the according prediction values
    # then we compute the log of the tensor to satisfy the cross entropy's need
    return -torch.log(y_hat.gather(1, y.view(-1, 1)))