Merge pull request #16 from VectorInstitute/b/reg-term-bugfix

B/reg term bugfix

Author: jwilles

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "vbll"
-version = "0.4.8"
+version = "0.4.9"
 description = ""
 authors = ["John Willes <johnwilles@gmail.com>"]
 readme = "README.md"

setup.py

@@ -3,7 +3,7 @@
 
 setup(
     name="vbll",
-    version="0.4.8",
+    version="0.4.9",
     packages=find_packages(),
     install_requires=["torch"],
 )

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_ll + kl_term_noise )
+            total_elbo -= self.regularization_weight * (kl_term_ll + kl_term_noise)
 
             return -total_elbo
 

vbll/layers/regression.py

@@ -404,7 +404,7 @@ def loss_fn(y):
             # compute expected KL
             kl_term_ll = torch.mean(grad_correction * expected_gaussian_kl(W, self.prior_scale, expect_sigma_inv))
             kl_term_noise = torch.mean(grad_correction * gaussian_kl(M, self.noise_prior_scale))
-            total_elbo -= self.regularization_weight * ( kl_term_ll + kl_term_noise )
+            total_elbo -= self.regularization_weight * (kl_term_ll + kl_term_noise)
             return -total_elbo
 
         return loss_fn