import torch


def cross_entropy_loss(stu_logits, tea_logits, temp=1.): # 使用交叉熵
  """ the same as nn.CrossEntropyLoss, but more flexible
  Args:
    stu_logits: tensor of shape [N, class]
    tea_logits: tensor of shape [N, class]
    temp: the distillation temperature
  Return:
    kd_loss: the cross entropy on soft labels
  """
  pred = torch.nn.functional.softmax(stu_logits, dim=1)
  labels = torch.nn.functional.softmax(tea_logits / temp, dim=1)
  kd_loss = (-labels * torch.log(pred)).sum(dim=1).mean()
  return kd_loss