def kl_divergence(self, logits_q, logits_p):
return (torch.exp(logits_q) * (logits_q - logits_p)).sum(<NUM_LIT:1>, keepdim=True)<EOL>
Categorical distribution KL divergence calculation
KL(Q || P) = sum Q_i log (Q_i / P_i)
When talking about logits this is:
sum exp(Q_i) * (Q_i - P_i)