diff --git a/examples/mnist/with_argbind_and_refactor.py b/examples/mnist/with_argbind_and_refactor.py
index ec009bd..c434dea 100644
--- a/examples/mnist/with_argbind_and_refactor.py
+++ b/examples/mnist/with_argbind_and_refactor.py
@@ -149,7 +149,7 @@ def scheduler(
     gamma : float, optional
         Learning rate step gamma, by default 0.7
     """
-    return StepLR(optimizer, step_size=1, gamma=gamma)
+    return StepLR(optimizer, step_size=step_size, gamma=gamma)
 
 @argbind.bind()
 def main(