with torch.no_grad():
      teacher_output = self.teacher(**inputs)

    # Compute soft targets for teacher and student
    soft_teacher = F.softmax(teacher_output.logits / self.temperature, dim=-1)
    soft_student = F.log_softmax(student_output.logits / self.temperature, dim=-1)

    # Compute the loss
    distillation_loss = self.loss_function(soft_student, soft_teacher) * (self.temperature ** 2)

    # Compute the true label loss
    student_target_loss = student_output.loss