Well this is my recon_loss function, and I’m using MSE.
def reconstruction_loss(x, x_recon, distribution):
   batch_size = x.size(0)
   assert batch_size != 0

if distribution == 'bernoulli':
    recon_loss = F.binary_cross_entropy_with_logits(x_recon, x, size_average=False).div(batch_size)
elif distribution == 'gaussian':
    x_recon = F.sigmoid(x_recon)
    recon_loss = F.mse_loss(x_recon, x, size_average=False).div(batch_size)
else:
    recon_loss = None

return recon_loss

Gaussian is being run in this case. But it’s not being called until after x_recon is output from my autoencoder. But my recon loss is being out put the decoder obviously during training I call:
x = Variable(cuda(x, self.use_cuda))
x_recon, mu, logvar = self.net(x)

This is based off almost exactly on: https://github.com/1Konny/Beta-VAE 1