So I tried to avoid spelling out the log likelihood but so there are two sets of mu and logvar:
def decode(self, z):
    h = F.relu(self.fc3(z))
    mu = self.fc4(h)
    logvar = self.fc5(h)
    return mu, logvar
 
def forward(self, x):
    mu_latent, logvar_latent = self.encode(x.view(-1, 784))
    z = self.reparameterize(mu, logvar)
    mu_x, logvar_x = self.decode(z)
    return mu_latent, logvar_latent, mu_x, logvar_x
 
def loss_function(x, mu_x, logvar_x, mu_latent, logvar_latent):
    # neg log likelihood of x under normal
    loss_rec = LOG_2_PI + logvar_x + (x - mu_x)**2 / (2*torch.exp(logvar_x))
    KLD = -0.5 * torch.sum(1 + logvar_latent - mu_latent.pow(2) - logvar_latent.exp())
    return loss_rec + KLD

I hope this approximately makes sense.
Best regards
Thomas