How to compute KL divergence for other distributions?

[Cero]

It seems like your formula is correct but one thread addresses your issue:

Using log_prob in distributions

Can someone please verify if the following use of log_prob fucntion is correct to calculate KL Divergence? self.q_y = RelaxedOneHotCategorical(self.temperature, logits=logits) self.log_q_y = self.q_y.log_prob(samples) self.p_y = RelaxedOneHotCategorical(self.temperature, logits=prior) self.log_p_y = self.p_y.log_prob(samples) kl = torch.exp(self.log_q_y) * (self.log_q_y - self.log_p_y)

It could be that you get wrong results because RelaxedOneHotCategorical is numerical instabil. The linked thread also provides a solution for that.