diff --git a/denoising_diffusion_pytorch/elucidated_diffusion.py b/denoising_diffusion_pytorch/elucidated_diffusion.py
index 2b3217e..740a1fe 100644
--- a/denoising_diffusion_pytorch/elucidated_diffusion.py
+++ b/denoising_diffusion_pytorch/elucidated_diffusion.py
@@ -140,16 +140,18 @@ class ElucidatedDiffusion(nn.Module):
     # sampling
 
     @torch.no_grad()
-    def sample(self, batch_size = 16):
+    def sample(self, batch_size = 16, num_sample_steps = None):
+        num_sample_steps = default(num_sample_steps, self.num_sample_steps)
+
         shape = (batch_size, self.channels, self.image_size, self.image_size)
 
         # get the schedule, which is returned as (sigma, gamma) tuple, and pair up with the next sigma and gamma
 
-        sigmas = self.sample_schedule()
+        sigmas = self.sample_schedule(num_sample_steps)
 
         gammas = torch.where(
             (sigmas >= self.S_tmin) & (sigmas <= self.S_tmax),
-            min(self.S_churn / self.num_sample_steps, sqrt(2) - 1),
+            min(self.S_churn / num_sample_steps, sqrt(2) - 1),
             0.
         )
 
diff --git a/setup.py b/setup.py
index 3c0d51a..ff1eb12 100644
--- a/setup.py
+++ b/setup.py
@@ -3,7 +3,7 @@ from setuptools import setup, find_packages
 setup(
   name = 'denoising-diffusion-pytorch',
   packages = find_packages(),
-  version = '0.22.0',
+  version = '0.23.0',
   license='MIT',
   description = 'Denoising Diffusion Probabilistic Models - Pytorch',
   author = 'Phil Wang',