diff --git a/transvae/loss.py b/transvae/loss.py
index 37c8fae..328dc65 100644
--- a/transvae/loss.py
+++ b/transvae/loss.py
@@ -11,7 +11,7 @@ def vae_loss(x, x_out, mu, logvar, weights, beta=1):
     x = x.contiguous().view(-1)
     x_out = x_out.contiguous().view(-1, x_out.size(2))
     BCE = F.cross_entropy(x_out, x, reduction='mean', weight=weights)
-    KLD = beta * -0.5 * torch.mean(1 + logvar - mu.pow(2) - logvar.exp())
+    KLD = beta * torch.mean(-0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp(), dim=1), dim=0)
     if torch.isnan(KLD):
         KLD = 0.
     return BCE + KLD, BCE, KLD
@@ -24,7 +24,7 @@ def trans_vae_loss(x, x_out, mu, logvar, true_len, pred_len, weights, beta=1):
     true_len = true_len.contiguous().view(-1)
     BCEmol = F.cross_entropy(x_out, x, reduction='mean', weight=weights)
     BCEmask = F.cross_entropy(pred_len, true_len, reduction='mean')
-    KLD = beta * -0.5 * torch.mean(1 + logvar - mu.pow(2) - logvar.exp())
+    KLD = beta * torch.mean(-0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp(), dim=1), dim=0)
     if torch.isnan(KLD):
         KLD = 0.
     return BCEmol + BCEmask + KLD, BCEmol, BCEmask, KLD