diff --git a/official/projects/waste_identification_ml/model_inference/Inference.ipynb b/official/projects/waste_identification_ml/model_inference/Inference.ipynb
index d5673d130b8..030b86a9aec 100644
--- a/official/projects/waste_identification_ml/model_inference/Inference.ipynb
+++ b/official/projects/waste_identification_ml/model_inference/Inference.ipynb
@@ -36,6 +36,17 @@
         "## Imports and Setup"
       ]
     },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "HhMR363skKaY"
+      },
+      "outputs": [],
+      "source": [
+        "!pip install -q tf_keras"
+      ]
+    },
     {
       "cell_type": "code",
       "execution_count": null,