make sure predicted mean is actually detached for all of the kl loss calculations

denoising_diffusion_pytorch/learned_gaussian_diffusion.py

@@ -130,10 +130,12 @@ class LearnedGaussianDiffusion(GaussianDiffusion):
 
         # kl loss with detached model predicted mean, for stability reasons as in paper
 
-        kl = normal_kl(true_mean, true_log_variance_clipped, model_mean.detach(), model_log_variance)
+        detached_model_mean = model_mean.detach()
+
+        kl = normal_kl(true_mean, true_log_variance_clipped, detached_model_mean, model_log_variance)
         kl = meanflat(kl) * NAT
 
-        decoder_nll = -discretized_gaussian_log_likelihood(x_start, means = model_mean, log_scales = 0.5 * model_log_variance)
+        decoder_nll = -discretized_gaussian_log_likelihood(x_start, means = detached_model_mean, log_scales = 0.5 * model_log_variance)
         decoder_nll = meanflat(decoder_nll) * NAT
 
         # at the first timestep return the decoder NLL, otherwise return KL(q(x_{t-1}|x_t,x_0) || p(x_{t-1}|x_t))

setup.py

@@ -3,7 +3,7 @@ from setuptools import setup, find_packages
 setup(
   name = 'denoising-diffusion-pytorch',
   packages = find_packages(),
-  version = '0.14.1',
+  version = '0.14.2',
   license='MIT',
   description = 'Denoising Diffusion Probabilistic Models - Pytorch',
   author = 'Phil Wang',