diff --git a/examples/mnist/mnist_e2e_on_prem.ipynb b/examples/mnist/mnist_e2e_on_prem.ipynb
index 7eb38b5b..e29fc9bb 100644
--- a/examples/mnist/mnist_e2e_on_prem.ipynb
+++ b/examples/mnist/mnist_e2e_on_prem.ipynb
@@ -132,6 +132,7 @@
     "from kubernetes import client as k8s_client\n",
     "from kubernetes import config as k8s_config\n",
     "from kubeflow.fairing.utils import is_running_in_k8s\n",
+    "import yaml\n",
     "\n",
     "pv_yaml = f'''\n",
     "apiVersion: v1\n",