Small error in the kl-divergence

[manuelhaussmann]

Unless I am missing something there is a slight error in the kl_divergence() definition in the _ConvNdGroupNJ class.

KLD_element = -self.weight_logvar + 0.5 * (self.weight_logvar.exp().pow(2) + self.weight_mu.pow(2)) - 0.5

treats self.weight_logvar as if it was the log std instead of the log variance. The correct expression should be (as in LinearGroupNJ):

KLD_element = -0.5 * self.weight_logvar + 0.5 * (self.weight_logvar.exp() + self.weight_mu.pow(2)) - 0.5

[clouizos]

Yes, you are indeed correct, thanks for pointing it out! We only tested the dense implementation so this fell through the cracks. We just pushed the fix.