diff --git a/sdk/python/jobs/single-step/tensorflow/mnist-distributed/tensorflow-mnist-distributed.ipynb b/sdk/python/jobs/single-step/tensorflow/mnist-distributed/tensorflow-mnist-distributed.ipynb
index 01f0989af6..a0ddaf8b34 100644
--- a/sdk/python/jobs/single-step/tensorflow/mnist-distributed/tensorflow-mnist-distributed.ipynb
+++ b/sdk/python/jobs/single-step/tensorflow/mnist-distributed/tensorflow-mnist-distributed.ipynb
@@ -118,10 +118,10 @@
    "cell_type": "code",
    "execution_count": null,
    "metadata": {
-    "name": "job",
     "gather": {
      "logged": 1634855420019
-    }
+    },
+    "name": "job"
    },
    "outputs": [],
    "source": [
@@ -134,12 +134,12 @@
     "    compute=\"cpu-cluster\",\n",
     "    instance_count=2,\n",
     "    # distribution = {\"type\": \"mpi\", \"process_count_per_instance\": 1},\n",
-    "    distribution={\n",
-    "        \"type\": \"tensorflow\",\n",
-    "        \"parameter_server_count\": 1,\n",
-    "        \"worker_count\": 2,\n",
-    "        \"added_property\": 7,\n",
-    "    },\n",
+    "    # distribution={\n",
+    "    #     \"type\": \"tensorflow\",\n",
+    "    #     \"parameter_server_count\": 1,  # for legacy TensorFlow 1.x\n",
+    "    #     \"worker_count\": 2,\n",
+    "    #     \"added_property\": 7,\n",
+    "    # },\n",
     "    # distribution = {\n",
     "    #        \"type\": \"pytorch\",\n",
     "    #        \"process_count_per_instance\": 4,\n",
@@ -151,7 +151,7 @@
     ")\n",
     "\n",
     "# can also set the distribution in a separate step and using the typed objects instead of a dict\n",
-    "job.distribution = TensorFlowDistribution(parameter_server_count=1, worker_count=2)"
+    "job.distribution = TensorFlowDistribution(worker_count=2)"
    ]
   },
   {
@@ -183,6 +183,16 @@
     "returned_job = ml_client.create_or_update(job)"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Wait until the job completes\n",
+    "ml_client.jobs.stream(returned_job.name)"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},