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

denoising_diffusion_pytorch/learned_gaussian_diffusion.py

@@ -107,6 +107,10 @@ class LearnedGaussianDiffusion(GaussianDiffusion):
         model_variance = model_log_variance.exp()
 
         x_start = self.predict_start_from_noise(x, t, pred_noise)
+
+        if clip_denoised:
+            x_start.clamp_(-1., 1.)
+
         model_mean, _, _ = self.q_posterior(x_start, x, t)
 
         return model_mean, model_variance, model_log_variance
@@ -126,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.0',
+  version = '0.14.2',
   license='MIT',
   description = 'Denoising Diffusion Probabilistic Models - Pytorch',
   author = 'Phil Wang',