diff --git a/2_mnist.ipynb b/2_mnist.ipynb
index ae71ed9..29e2bdc 100644
--- a/2_mnist.ipynb
+++ b/2_mnist.ipynb
@@ -79,8 +79,8 @@
    },
    "outputs": [],
    "source": [
-    "def model_fn(x, y, mode)\n",
-    "    logits = conv_model(x, mode == tf.config.learn.ModeKeys.TRAIN)\n",
+    "def model_fn(x, y, mode):\n",
+    "    logits = conv_model(x, mode == tf.contrib.learn.ModeKeys.TRAIN)\n",
     "    if mode != tf.contrib.learn.ModeKeys.INFER:\n",
     "        loss = tf.contrib.losses.softmax_cross_entropy(logits, y)\n",
     "        if mode == tf.contrib.learn.ModeKeys.TRAIN:\n",