diff --git a/tests/notebooks/gpu/kfp-tensorflow/kfp-tensorflow-integration.ipynb b/tests/notebooks/gpu/kfp-tensorflow/kfp-tensorflow-integration.ipynb
index 8c77b05..081db64 100644
--- a/tests/notebooks/gpu/kfp-tensorflow/kfp-tensorflow-integration.ipynb
+++ b/tests/notebooks/gpu/kfp-tensorflow/kfp-tensorflow-integration.ipynb
@@ -67,10 +67,10 @@
     "    HTTPS_PROXY = os.environ['HTTPS_PROXY']\n",
     "    NO_PROXY = os.environ['NO_PROXY']\n",
     "\n",
-    "def add_proxy(obj, http_proxy=HTTP_PROXY, https_proxy=HTTPS_PROXY, no_proxy=NO_PROXY):\n",
+    "def add_proxy(task: dsl.PipelineTask, http_proxy=HTTP_PROXY, https_proxy=HTTPS_PROXY, no_proxy=NO_PROXY) -> dsl.PipelineTask:\n",
     "    \"\"\"Adds the proxy env vars to the PipelineTask object.\"\"\"\n",
     "    return (\n",
-    "        obj.set_env_variable(name='http_proxy', value=http_proxy)\n",
+    "        task.set_env_variable(name='http_proxy', value=http_proxy)\n",
     "        .set_env_variable(name='https_proxy', value=https_proxy)\n",
     "        .set_env_variable(name='HTTP_PROXY', value=http_proxy)\n",
     "        .set_env_variable(name='HTTPS_PROXY', value=https_proxy)\n",
@@ -102,9 +102,9 @@
     "        raise RuntimeError(\"No GPU has been detected.\")\n",
     "    return str(len(gpus)>0)\n",
     "\n",
-    "def add_gpu_request(obj):\n",
+    "def add_gpu_request(task: dsl.PipelineTask) -> dsl.PipelineTask:\n",
     "    \"\"\"Add a request field for a GPU to the container created by the PipelineTask object.\"\"\"\n",
-    "    return ( obj.add_node_selector_constraint(accelerator = \"nvidia.com/gpu\").set_accelerator_limit(limit = 1) )"
+    "    return ( task.add_node_selector_constraint(accelerator = \"nvidia.com/gpu\").set_accelerator_limit(limit = 1) )"
    ]
   },
   {