fix vamp llh calculation (#27)

clementchadebec · Jul 1, 2022 · 1588bac · 1588bac
1 parent b0238cf
commit 1588bac
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/src/pythae/models/vamp/vamp_model.py b/src/pythae/models/vamp/vamp_model.py
@@ -155,9 +155,10 @@ def _log_p_z(self, z):
 
         log_p_z = (
             torch.sum(
-                -0.5
-                * (prior_log_var + (z_expand - prior_mu) ** 2)
-                / prior_log_var.exp(),
+                -0.5 * (
+                    prior_log_var + (z_expand - prior_mu) ** 2
+                    / torch.exp(prior_log_var)
+                ),
                 dim=2,
             )
             - torch.log(torch.tensor(C).type(torch.float))