import torch


def soft_cross_entropy_tinybert(input, targets):
    """ Soft Cross Entropy loss for hinton's dark knowledge

        Args:
            input (`Tensor`): shape of [None, N]
            targets (`Tensor`): shape of [None, N]
        Returns:
            loss (`Tensor`): scalar tensor
    """
    student_likelihood = torch.nn.functional.log_softmax(input, dim=-1)
    targets_prob = torch.nn.functional.softmax(targets, dim=-1)
    return (- targets_prob * student_likelihood).mean()