diff --git a/site/en/tutorials/keras/text_classification_with_hub.ipynb b/site/en/tutorials/keras/text_classification_with_hub.ipynb
index b43e59e35f..e7e4288ef3 100644
--- a/site/en/tutorials/keras/text_classification_with_hub.ipynb
+++ b/site/en/tutorials/keras/text_classification_with_hub.ipynb
@@ -121,7 +121,8 @@
       "outputs": [],
       "source": [
         "!pip install tensorflow-hub\n",
-        "!pip install tensorflow-datasets"
+        "!pip install tensorflow-datasets\n",
+        "!pip install tf-keras"
       ]
     },
     {
@@ -138,6 +139,7 @@
         "import tensorflow as tf\n",
         "import tensorflow_hub as hub\n",
         "import tensorflow_datasets as tfds\n",
+        "import tf_keras as keras\n",
         "\n",
         "print(\"Version: \", tf.__version__)\n",
         "print(\"Eager mode: \", tf.executing_eagerly())\n",
@@ -290,10 +292,10 @@
       },
       "outputs": [],
       "source": [
-        "model = tf.keras.Sequential()\n",
+        "model = keras.Sequential()\n",
         "model.add(hub_layer)\n",
-        "model.add(tf.keras.layers.Dense(16, activation='relu'))\n",
-        "model.add(tf.keras.layers.Dense(1))\n",
+        "model.add(keras.layers.Dense(16, activation='relu'))\n",
+        "model.add(keras.layers.Dense(1))\n",
         "\n",
         "model.summary()"
       ]
@@ -339,7 +341,7 @@
       "outputs": [],
       "source": [
         "model.compile(optimizer='adam',\n",
-        "              loss=tf.keras.losses.BinaryCrossentropy(from_logits=True),\n",
+        "              loss=keras.losses.BinaryCrossentropy(from_logits=True),\n",
         "              metrics=['accuracy'])"
       ]
     },