diff --git a/tensorflow_privacy/privacy/membership_inference_attack/codelab.ipynb b/tensorflow_privacy/privacy/membership_inference_attack/codelab.ipynb
index ad13f10..ff78a4e 100644
--- a/tensorflow_privacy/privacy/membership_inference_attack/codelab.ipynb
+++ b/tensorflow_privacy/privacy/membership_inference_attack/codelab.ipynb
@@ -135,6 +135,16 @@
         "from tensorflow_privacy.privacy.membership_inference_attack import membership_inference_attack as mia"
       ]
     },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "colab_type": "text",
+        "id": "pBbcG86th_sW"
+      },
+      "source": [
+        "## Train a simple model on CIFAR10 with Keras."
+      ]
+    },
     {
       "cell_type": "code",
       "execution_count": null,
@@ -145,8 +155,6 @@
       },
       "outputs": [],
       "source": [
-        "#@title Train a simple model on CIFAR10 with Keras.\n",
-        "\n",
         "dataset = 'cifar10'\n",
         "num_classes = 10\n",
         "num_conv = 3\n",
@@ -254,6 +262,16 @@
         "loss_test = cce(constant(y_test), constant(prob_test), from_logits=False).numpy()"
       ]
     },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "colab_type": "text",
+        "id": "QETxVOHLiHP4"
+      },
+      "source": [
+        "## Run membership inference attacks."
+      ]
+    },
     {
       "cell_type": "code",
       "execution_count": null,
@@ -264,7 +282,6 @@
       },
       "outputs": [],
       "source": [
-        "#@title Run membership inference attacks.\n",
         "#@markdown We will now execute membership inference attack against the\n",
         "#@markdown previously trained CIFAR10 model. This will generate a number of\n",
         "#@markdown scores (most notably, attacker advantage and AUC for the membership\n",