diff --git a/TensorRT/demo.py b/TensorRT/demo.py
index 456fad0..c856ea4 100644
--- a/TensorRT/demo.py
+++ b/TensorRT/demo.py
@@ -110,7 +110,7 @@ def predict(
                 )
 
                 # compute the previous noisy sample x_t -> x_t-1
-                latents = self.scheduler.step(noise_pred.cuda(), t, latents)[
+                latents = self.scheduler.step(noise_pred.cuda(), i, latents)[
                     "prev_sample"
                 ]