diff --git a/examples/research_projects/onnxruntime/unconditional_image_generation/train_unconditional.py b/examples/research_projects/onnxruntime/unconditional_image_generation/train_unconditional.py
index b9dec9a745e8..586ed8c8d397 100644
--- a/examples/research_projects/onnxruntime/unconditional_image_generation/train_unconditional.py
+++ b/examples/research_projects/onnxruntime/unconditional_image_generation/train_unconditional.py
@@ -552,6 +552,7 @@ def transform_images(examples):
                     generator=generator,
                     batch_size=args.eval_batch_size,
                     output_type="numpy",
+                    num_inference_steps=args.ddpm_num_steps,
                 ).images
 
                 # denormalize the images and save to tensorboard
diff --git a/examples/unconditional_image_generation/train_unconditional.py b/examples/unconditional_image_generation/train_unconditional.py
index 324ddb0538e3..32dd96c88249 100644
--- a/examples/unconditional_image_generation/train_unconditional.py
+++ b/examples/unconditional_image_generation/train_unconditional.py
@@ -587,6 +587,7 @@ def transform_images(examples):
                     generator=generator,
                     batch_size=args.eval_batch_size,
                     output_type="numpy",
+                    num_inference_steps=args.ddpm_num_steps,
                 ).images
 
                 # denormalize the images and save to tensorboard