clamp the sigma coming out of the log normal distribution

andrewmoise · Jun 28, 2022 · 6184937 · 6184937
1 parent 76b79aa
commit 6184937
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/denoising_diffusion_pytorch/elucidated_diffusion.py b/denoising_diffusion_pytorch/elucidated_diffusion.py
@@ -99,7 +99,8 @@ def c_noise(self, sigma):
     # noise distribution
 
     def noise_distribution(self, batch_size):
-        return (self.P_mean + self.P_std * torch.randn((batch_size,), device = self.device)).exp()
+        sigmas = (self.P_mean + self.P_std * torch.randn((batch_size,), device = self.device)).exp()
+        return sigmas.clamp(min = self.sigma_min, max =self.sigma_max)
 
     def loss_weight(self, sigma):
         return (sigma ** 2 + self.sigma_data ** 2) * (sigma * self.sigma_data) ** -2