diff --git a/tensorflow/contrib/opt/python/training/weight_decay_optimizers.py b/tensorflow/contrib/opt/python/training/weight_decay_optimizers.py
index 1e8351b70ff..8b8065c678e 100644
--- a/tensorflow/contrib/opt/python/training/weight_decay_optimizers.py
+++ b/tensorflow/contrib/opt/python/training/weight_decay_optimizers.py
@@ -64,10 +64,10 @@ class DecoupledWeightDecayExtension(object):
   the decay to the `weight_decay` as well. For example:
 
   ```python
-    decay = tf.train.piecewise_constant(tf.train.get_global_step(), 
-                                        [10000, 15000], [1e-1, 1e-2, 1e-3])
-    lr = 1*decay
-    wd = 1e-4*decay
+    schedule = tf.train.piecewise_constant(tf.train.get_global_step(), 
+                                           [10000, 15000], [1e-0, 1e-1, 1e-2])
+    lr = 1e-1 * schedule()
+    wd = lambda: 1e-4 * schedule()
 
     # ...