updating KLD function

dev4488 · Sep 4, 2018 · 798d481 · 798d481
1 parent dd931d8
commit 798d481
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/vae_gumbel_softmax.py b/vae_gumbel_softmax.py
@@ -126,7 +126,7 @@ def loss_function(recon_x, x,qy):
 
     log_qy = torch.log(qy+1e-20)
     g = Variable(torch.log(torch.Tensor([1.0/categorical_dim])).cuda())
-    KLD = -torch.sum(qy*(log_qy - g))
+    KLD = torch.sum(qy*(log_qy - g),dim=-1).mean()
 
     return BCE + KLD