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 f681ee2087f..a8c3b04e003 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
@@ -130,7 +130,7 @@
     "    code=\"./src\",  # local path where the code is stored\n",
     "    command=\"python main.py --epochs ${{inputs.epochs}} --model-dir ${{inputs.model_dir}}\",\n",
     "    inputs={\"epochs\": 1, \"model_dir\": \"outputs/keras-model\"},\n",
-    "    environment=\"AzureML-tensorflow-2.12-cuda11@latest\",\n",
+    "    environment=\"AzureML-tensorflow-2.16-cuda11@latest\",\n",
     "    compute=\"cpu-cluster\",\n",
     "    instance_count=2,\n",
     "    # distribution = {\"type\": \"mpi\", \"process_count_per_instance\": 1},\n",