Bug fix on KL weight for heteroscedastic classification

Author: jmesharrison

vbll/layers/classification.py

@@ -525,7 +525,7 @@ def loss_fn(y):
             kl_term_noise = gaussian_kl(self.M, self.noise_prior_scale)
 
             total_elbo = torch.mean(self.softmax_bound(x, y))
-            total_elbo -= self.regularization_weight * (kl_term_noise + kl_term_ll)
+            total_elbo -= self.regularization_weight * kl_term_noise + kl_term_ll
 
             return -total_elbo