typo in KLDivTeacherList?

[jihyukkim-nlp]

Thank you for sharing codes. I wonder whether there is a typo in the implementation of KLDivTeacherList class.

The implementation is

class KLDivTeacherList(nn.Module):
    def __init__(self):
        super(KLDivTeacherList, self).__init__()
        self.kl = torch.nn.KLDivLoss(reduction="batchmean")
    def forward(self, scores, labels):
        loss = self.kl(scores.softmax(-1),labels.softmax(-1)) # is this a typo?
        return loss

However, PyTorch documentation for KLDivLoss (https://pytorch.org/docs/1.9.1/generated/torch.nn.KLDivLoss.html) says

• As with NLLLoss, the input given is expected to contain log-probabilities.
• The targets are interpreted as probabilities by default,

So, from what I understand, the forward function should be

def forward(self, scores, labels):
    # loss = self.kl(scores.softmax(-1),labels.softmax(-1)) # is this a typo?
    # before: softmax of scores

    # after : log-softmax of scores
    loss = self.kl(torch.nn.functional.log_softmax(scores, dim=-1), torch.nn.functional.softmax(labels, dim=-1))
    return loss

I wonder if this is a typo or if I'm missing something. Thanks in advance.