diff --git a/how-to-use-azureml/work-with-data/datasets-tutorial/pipeline-with-datasets/keras-mnist-fashion/train.py b/how-to-use-azureml/work-with-data/datasets-tutorial/pipeline-with-datasets/keras-mnist-fashion/train.py
index b0215ad1..7cb8d72a 100644
--- a/how-to-use-azureml/work-with-data/datasets-tutorial/pipeline-with-datasets/keras-mnist-fashion/train.py
+++ b/how-to-use-azureml/work-with-data/datasets-tutorial/pipeline-with-datasets/keras-mnist-fashion/train.py
@@ -1,10 +1,11 @@
-import keras
-from keras.models import Sequential
-from keras.layers import Dense, Dropout, Flatten
-from keras.layers import Conv2D, MaxPooling2D
-from keras.layers.normalization import BatchNormalization
-from keras.utils import to_categorical
-from keras.callbacks import Callback
+from tensorflow.keras.models import Sequential
+from tensorflow.keras.layers import Dense, Dropout, Flatten
+from tensorflow.keras.layers import Conv2D, MaxPooling2D
+from tensorflow.keras.layers import BatchNormalization
+from tensorflow.keras.losses import categorical_crossentropy
+from tensorflow.keras.optimizers import Adam
+from tensorflow.keras.utils import to_categorical
+from tensorflow.keras.callbacks import Callback
 
 import numpy as np
 import pandas as pd
@@ -64,8 +65,8 @@ model.add(Dense(128, activation='relu'))
 model.add(Dropout(0.3))
 model.add(Dense(num_classes, activation='softmax'))
 
-model.compile(loss=keras.losses.categorical_crossentropy,
-              optimizer=keras.optimizers.Adam(),
+model.compile(loss=categorical_crossentropy,
+              optimizer=Adam(),
               metrics=['accuracy'])
 
 # start an Azure ML run
diff --git a/how-to-use-azureml/work-with-data/datasets-tutorial/pipeline-with-datasets/pipeline-for-image-classification.ipynb b/how-to-use-azureml/work-with-data/datasets-tutorial/pipeline-with-datasets/pipeline-for-image-classification.ipynb
index 648dfce4..14044ce4 100644
--- a/how-to-use-azureml/work-with-data/datasets-tutorial/pipeline-with-datasets/pipeline-for-image-classification.ipynb
+++ b/how-to-use-azureml/work-with-data/datasets-tutorial/pipeline-with-datasets/pipeline-for-image-classification.ipynb
@@ -270,16 +270,19 @@
         "%%writefile conda_dependencies.yml\n",
         "\n",
         "dependencies:\n",
-        "- python=3.6.2\n",
+        "- python=3.8\n",
+        "- pip==20.2.4\n",
         "- pip:\n",
         "  - azureml-core\n",
         "  - azureml-dataset-runtime\n",
-        "  - keras==2.4.3\n",
-        "  - tensorflow==2.4.3\n",
+        "  - keras==2.6\n",
+        "  - tensorflow-gpu==2.6\n",
         "  - numpy\n",
         "  - scikit-learn\n",
         "  - pandas\n",
-        "  - matplotlib"
+        "  - matplotlib\n",
+        "  - protobuf==3.20.1\n",
+        "  - typing-extensions==4.3.0"
       ]
     },
     {