fix off by one

andrewmoise · Jun 28, 2022 · b87ea27 · b87ea27
1 parent a8403b8
commit b87ea27
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/denoising_diffusion_pytorch/elucidated_diffusion.py b/denoising_diffusion_pytorch/elucidated_diffusion.py
@@ -116,7 +116,7 @@ def sample_schedule(self, num_sample_steps = None):
         N = num_sample_steps
         inv_rho = 1 / rho
 
-        for i in range(num_sample_steps - 1):
+        for i in range(num_sample_steps):
             sigma_i = (sigma_max ** inv_rho + i / (N - 1) * (sigma_min ** inv_rho - sigma_max ** inv_rho)) ** rho
             gamma_i = gamma if S_tmin <= sigma_i <= S_tmax else 0.
             yield sigma_i, gamma_i