diff --git a/notebooks/official/ml_metadata/vertex-pipelines-ml-metadata.ipynb b/notebooks/official/ml_metadata/vertex-pipelines-ml-metadata.ipynb
index eb1b07097..4ced1547b 100644
--- a/notebooks/official/ml_metadata/vertex-pipelines-ml-metadata.ipynb
+++ b/notebooks/official/ml_metadata/vertex-pipelines-ml-metadata.ipynb
@@ -143,8 +143,8 @@
       },
       "outputs": [],
       "source": [
-        "! pip3 install --upgrade --quiet google-cloud-aiplatform==1.7.0 \\\n",
-        "                                 kfp==1.8.9"
+        "! pip3 install --upgrade --quiet google-cloud-aiplatform \\\n",
+        "                                 'kfp<2.0'"
       ]
     },
     {
@@ -379,6 +379,9 @@
       },
       "outputs": [],
       "source": [
+        "import sys\n",
+        "\n",
+        "IS_COLAB = \"google.colab\" in sys.modules\n",
         "if (\n",
         "    SERVICE_ACCOUNT == \"\"\n",
         "    or SERVICE_ACCOUNT is None\n",
@@ -1121,6 +1124,27 @@
         "\n",
         "* Delete the [Cloud Storage bucket](https://console.cloud.google.com/storage/browser/) you created"
       ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "f2c21373a498"
+      },
+      "outputs": [],
+      "source": [
+        "import os\n",
+        "\n",
+        "try:\n",
+        "    run1.delete()\n",
+        "    run2.delete()\n",
+        "except Exception as e:\n",
+        "    print(e)\n",
+        "\n",
+        "delete_bucket = False\n",
+        "if delete_bucket or os.getenv(\"IS_TESTING\"):\n",
+        "    ! gsutil rm -rf {BUCKET_URI}"
+      ]
     }
   ],
   "metadata": {