diff --git a/aws_sagemaker_studio/frameworks/tensorflow_mnist/tensorflow_mnist.ipynb b/aws_sagemaker_studio/frameworks/tensorflow_mnist/tensorflow_mnist.ipynb
index f5372c44a4..2047b9bade 100644
--- a/aws_sagemaker_studio/frameworks/tensorflow_mnist/tensorflow_mnist.ipynb
+++ b/aws_sagemaker_studio/frameworks/tensorflow_mnist/tensorflow_mnist.ipynb
@@ -137,11 +137,10 @@
     "\n",
     "estimator = TensorFlow(entry_point=training_script,\n",
     "                       role=role,\n",
-    "                       train_instance_count=1,\n",
-    "                       train_instance_type='ml.p2.xlarge',\n",
+    "                       instance_count=1,\n",
+    "                       instance_type='ml.p2.xlarge',\n",
     "                       framework_version=tf_version,\n",
-    "                       py_version='py3',\n",
-    "                       distributions={'parameter_server': {'enabled': True}})"
+    "                       py_version='py3')"
    ]
   },
   {
@@ -259,6 +258,13 @@
    "source": [
     "predictor.delete_endpoint()"
    ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
   }
  ],
  "metadata": {
@@ -266,7 +272,7 @@
   "kernelspec": {
    "display_name": "Python 3 (Data Science)",
    "language": "python",
-   "name": "python3__SAGEMAKER_INTERNAL__arn:aws:sagemaker:us-east-2:429704687514:image/datascience-1.0"
+   "name": "python3__SAGEMAKER_INTERNAL__arn:aws:sagemaker:us-east-1:081325390199:image/datascience-1.0"
   },
   "language_info": {
    "codemirror_mode": {