diff --git a/cli/endpoints/batch/mnist/code/batch_driver.py b/cli/endpoints/batch/mnist/code/batch_driver.py
index 3d5a1113d8..e914ede6cc 100644
--- a/cli/endpoints/batch/mnist/code/batch_driver.py
+++ b/cli/endpoints/batch/mnist/code/batch_driver.py
@@ -4,7 +4,7 @@
 import os
 import numpy as np
 import pandas as pd
-import tensorflow as tf
+import tensorflow.compat.v1 as tf
 from PIL import Image
 from azureml.core import Model
 
@@ -18,6 +18,7 @@ def init():
     model_path = os.path.join(os.environ["AZUREML_MODEL_DIR"], "model")
 
     # contruct graph to execute
+    tf.disable_v2_behavior()
     tf.reset_default_graph()
     saver = tf.train.import_meta_graph(os.path.join(model_path, "mnist-tf.model.meta"))
     g_tf_sess = tf.Session(config=tf.ConfigProto(device_count={"GPU": 0}))
diff --git a/cli/endpoints/batch/mnist/environment/conda.json b/cli/endpoints/batch/mnist/environment/conda.json
index ee8535d49d..0397e7d532 100644
--- a/cli/endpoints/batch/mnist/environment/conda.json
+++ b/cli/endpoints/batch/mnist/environment/conda.json
@@ -4,11 +4,11 @@
     "conda-forge"
   ],
   "dependencies": [
-    "python=3.6.2",
+    "python=3.9.16",
     "pip",
     {
       "pip": [
-        "tensorflow==2.1.4",
+        "tensorflow==2.11.0",
         "pillow",
         "azureml-core",
         "azureml-dataset-runtime[fuse]"
diff --git a/cli/endpoints/batch/mnist/environment/conda.yml b/cli/endpoints/batch/mnist/environment/conda.yml
index 8675e97999..9908c0b718 100644
--- a/cli/endpoints/batch/mnist/environment/conda.yml
+++ b/cli/endpoints/batch/mnist/environment/conda.yml
@@ -2,10 +2,10 @@ name: mnist-env
 channels:
   - conda-forge
 dependencies:
-  - python=3.6.2
+  - python=3.9.16
   - pip<22.0
   - pip:
-    - tensorflow==2.1.4
+    - tensorflow==2.11.0
     - pillow
     - pandas
     - azureml-core