diff --git a/site/en/guide/autodiff.ipynb b/site/en/guide/autodiff.ipynb
index 237a224569..e07f4450a1 100644
--- a/site/en/guide/autodiff.ipynb
+++ b/site/en/guide/autodiff.ipynb
@@ -799,12 +799,14 @@
       "source": [
         "x = tf.Variable(2.0)\n",
         "\n",
-        "for epoch in range(2):\n",
+        "for epoch in range(3):\n",
         "  with tf.GradientTape() as tape:\n",
-        "    y = x+1\n",
+        "    y = x**2\n",
         "\n",
+        "  print("epoch:", epoch)\n,
         "  print(type(x).__name__, \":\", tape.gradient(y, x))\n",
-        "  x = x + 1   # This should be `x.assign_add(1)`"
+        "  x = x + 1   # The `tf.Variable` has been inadvertently replaced with a `tf.Tensor`."
+        "              # This should be `x.assign_add(1)`."
       ]
     },
     {