examples/research_projects/zero-shot-distillation/distill_classifier.py/DistillationTrainer/compute_loss
class DistillationTrainer:    def compute_loss(self, model, inputs, return_outputs=False):
        target_p = inputs["labels"]
        outputs = model(inputs["input_ids"], attention_mask=inputs["attention_mask"])
        logits = outputs[0]

        loss = -torch.sum(target_p * logits.log_softmax(dim=-1), axis=-1).mean()

        if return_outputs:
            return loss, outputs

        return loss