diff --git a/tensorflow/lite/g3doc/performance/post_training_float16_quant.ipynb b/tensorflow/lite/g3doc/performance/post_training_float16_quant.ipynb
index cf589a2b968..ef08902865e 100644
--- a/tensorflow/lite/g3doc/performance/post_training_float16_quant.ipynb
+++ b/tensorflow/lite/g3doc/performance/post_training_float16_quant.ipynb
@@ -61,6 +61,9 @@
         "  \u003ctd\u003e\n",
         "    \u003ca target=\"_blank\" href=\"https://github.com/tensorflow/tensorflow/blob/master/tensorflow/lite/g3doc/performance/post_training_float16_quant.ipynb\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" /\u003eView source on GitHub\u003c/a\u003e\n",
         "  \u003c/td\u003e\n",
+        "  \u003ctd\u003e\n",
+        "    \u003ca href=\"https://storage.googleapis.com/tensorflow_docs/tensorflow/lite/g3doc/performance/post_training_float16_quant.ipynb\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/download_logo_32px.png\" /\u003eDownload notebook\u003c/a\u003e\n",
+        "  \u003c/td\u003e\n",
         "\u003c/table\u003e"
       ]
     },
@@ -113,13 +116,7 @@
         "import logging\n",
         "logging.getLogger(\"tensorflow\").setLevel(logging.DEBUG)\n",
         "\n",
-        "try:\n",
-        "  # %tensorflow_version only exists in Colab.\n",
-        "  import tensorflow.compat.v2 as tf\n",
-        "except Exception:\n",
-        "  pass\n",
-        "tf.enable_v2_behavior()\n",
-        "\n",
+        "import tensorflow as tf\n",
         "from tensorflow import keras\n",
         "import numpy as np\n",
         "import pathlib"
@@ -173,12 +170,12 @@
         "  keras.layers.Conv2D(filters=12, kernel_size=(3, 3), activation=tf.nn.relu),\n",
         "  keras.layers.MaxPooling2D(pool_size=(2, 2)),\n",
         "  keras.layers.Flatten(),\n",
-        "  keras.layers.Dense(10, activation=tf.nn.softmax)\n",
+        "  keras.layers.Dense(10)\n",
         "])\n",
         "\n",
         "# Train the digit classification model\n",
         "model.compile(optimizer='adam',\n",
-        "              loss='sparse_categorical_crossentropy',\n",
+        "              loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),\n",
         "              metrics=['accuracy'])\n",
         "model.fit(\n",
         "  train_images,\n",
@@ -597,28 +594,11 @@
         "\n",
         "Detailed documentation on the TFLite GPU delegate and how to use it in your application can be found [here](https://www.tensorflow.org/lite/performance/gpu_advanced?source=post_page---------------------------)"
       ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": 0,
-      "metadata": {
-        "colab": {},
-        "colab_type": "code",
-        "id": "BeUSdwKVixvk"
-      },
-      "outputs": [],
-      "source": [
-        ""
-      ]
     }
   ],
   "metadata": {
     "colab": {
       "collapsed_sections": [],
-      "last_runtime": {
-        "build_target": "//learning/brain/python/client:colab_notebook_py3",
-        "kind": "private"
-      },
       "name": "post_training-float16-quant.ipynb",
       "private_outputs": true,
       "provenance": [],
diff --git a/tensorflow/lite/g3doc/performance/post_training_integer_quant.ipynb b/tensorflow/lite/g3doc/performance/post_training_integer_quant.ipynb
index fddee15bc1d..ad461f56d6f 100644
--- a/tensorflow/lite/g3doc/performance/post_training_integer_quant.ipynb
+++ b/tensorflow/lite/g3doc/performance/post_training_integer_quant.ipynb
@@ -61,6 +61,9 @@
         "  \u003ctd\u003e\n",
         "    \u003ca target=\"_blank\" href=\"https://github.com/tensorflow/tensorflow/blob/master/tensorflow/lite/g3doc/performance/post_training_integer_quant.ipynb\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" /\u003eView source on GitHub\u003c/a\u003e\n",
         "  \u003c/td\u003e\n",
+        "  \u003ctd\u003e\n",
+        "    \u003ca href=\"https://storage.googleapis.com/tensorflow_docs/tensorflow/lite/g3doc/performance/post_training_integer_quant.ipynb\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/download_logo_32px.png\" /\u003eDownload notebook\u003c/a\u003e\n",
+        "  \u003c/td\u003e\n",
         "\u003c/table\u003e"
       ]
     },
@@ -114,13 +117,7 @@
         "import logging\n",
         "logging.getLogger(\"tensorflow\").setLevel(logging.DEBUG)\n",
         "\n",
-        "try:\n",
-        "  # %tensorflow_version only exists in Colab.\n",
-        "  import tensorflow.compat.v2 as tf\n",
-        "except Exception:\n",
-        "  pass\n",
-        "tf.enable_v2_behavior()\n",
-        "\n",
+        "import tensorflow as tf\n",
         "from tensorflow import keras\n",
         "import numpy as np\n",
         "import pathlib"
@@ -158,15 +155,15 @@
         "model = keras.Sequential([\n",
         "  keras.layers.InputLayer(input_shape=(28, 28)),\n",
         "  keras.layers.Reshape(target_shape=(28, 28, 1)),\n",
-        "  keras.layers.Conv2D(filters=12, kernel_size=(3, 3), activation=tf.nn.relu),\n",
+        "  keras.layers.Conv2D(filters=12, kernel_size=(3, 3), activation='relu'),\n",
         "  keras.layers.MaxPooling2D(pool_size=(2, 2)),\n",
         "  keras.layers.Flatten(),\n",
-        "  keras.layers.Dense(10, activation=tf.nn.softmax)\n",
+        "  keras.layers.Dense(10)\n",
         "])\n",
         "\n",
         "# Train the digit classification model\n",
         "model.compile(optimizer='adam',\n",
-        "              loss='sparse_categorical_crossentropy',\n",
+        "              loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),\n",
         "              metrics=['accuracy'])\n",
         "model.fit(\n",
         "  train_images,\n",
@@ -277,7 +274,7 @@
       },
       "outputs": [],
       "source": [
-        "converter.optimizations = [tf.lite.Optimize.OPTIMIZE_FOR_SIZE]"
+        "converter.optimizations = [tf.lite.Optimize.DEFAULT]"
       ]
     },
     {
@@ -655,10 +652,6 @@
   "metadata": {
     "colab": {
       "collapsed_sections": [],
-      "last_runtime": {
-        "build_target": "//learning/brain/python/client:colab_notebook_py3",
-        "kind": "private"
-      },
       "name": "post_training_integer_quant.ipynb",
       "private_outputs": true,
       "provenance": [],
diff --git a/tensorflow/lite/g3doc/performance/post_training_quant.ipynb b/tensorflow/lite/g3doc/performance/post_training_quant.ipynb
index d6edb656d0e..201ccf5bdc3 100644
--- a/tensorflow/lite/g3doc/performance/post_training_quant.ipynb
+++ b/tensorflow/lite/g3doc/performance/post_training_quant.ipynb
@@ -61,6 +61,9 @@
         "  \u003ctd\u003e\n",
         "    \u003ca target=\"_blank\" href=\"https://github.com/tensorflow/tensorflow/blob/master/tensorflow/lite/g3doc/performance/post_training_quant.ipynb\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" /\u003eView source on GitHub\u003c/a\u003e\n",
         "  \u003c/td\u003e\n",
+        "  \u003ctd\u003e\n",
+        "    \u003ca href=\"https://storage.googleapis.com/tensorflow_docs/tensorflow/lite/g3doc/performance/post_training_quant.ipynb\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/download_logo_32px.png\" /\u003eDownload notebook\u003c/a\u003e\n",
+        "  \u003c/td\u003e\n",
         "\u003c/table\u003e"
       ]
     },
@@ -134,13 +137,7 @@
         "import logging\n",
         "logging.getLogger(\"tensorflow\").setLevel(logging.DEBUG)\n",
         "\n",
-        "try:\n",
-        "  # %tensorflow_version only exists in Colab.\n",
-        "  import tensorflow.compat.v2 as tf\n",
-        "except Exception:\n",
-        "  pass\n",
-        "tf.enable_v2_behavior()\n",
-        "\n",
+        "import tensorflow as tf\n",
         "from tensorflow import keras\n",
         "import numpy as np\n",
         "import pathlib"
@@ -181,12 +178,12 @@
         "  keras.layers.Conv2D(filters=12, kernel_size=(3, 3), activation=tf.nn.relu),\n",
         "  keras.layers.MaxPooling2D(pool_size=(2, 2)),\n",
         "  keras.layers.Flatten(),\n",
-        "  keras.layers.Dense(10, activation=tf.nn.softmax)\n",
+        "  keras.layers.Dense(10)\n",
         "])\n",
         "\n",
         "# Train the digit classification model\n",
         "model.compile(optimizer='adam',\n",
-        "              loss='sparse_categorical_crossentropy',\n",
+        "              loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),\n",
         "              metrics=['accuracy'])\n",
         "model.fit(\n",
         "  train_images,\n",
@@ -620,10 +617,6 @@
   "metadata": {
     "colab": {
       "collapsed_sections": [],
-      "last_runtime": {
-        "build_target": "//learning/brain/python/client:colab_notebook_py3",
-        "kind": "private"
-      },
       "name": "post_training_quant.ipynb",
       "private_outputs": true,
       "provenance": [],