Is the mean value calculated using these lines of code?
    logpx_z=-torch.sum(BCE,[1,2],keepdim=False)
    logpz=log_normal_pdf(z2,torch.tensor(0.),torch.tensor(0.))
    logqz_x=log_normal_pdf(z2, mean, logvar)
    mean=logpx_z+logpz-logqz_x

If so, could you try to store the activations and upload them so that we could try to reproduce this issue?