From 917eb8e2a03919cf9a06e5adffb77d9f4c18f390 Mon Sep 17 00:00:00 2001
From: Vikas Kumar <vikumar@amazon.com>
Date: Thu, 5 Nov 2020 13:48:50 -0800
Subject: [PATCH] pin tensorflow dataset in test config

---
 config/buildspec_vanilla_framework_tests.yml | 2 +-
 config/install_smdebug.sh                    | 2 +-
 config/tests.sh                              | 2 +-
 3 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/config/buildspec_vanilla_framework_tests.yml b/config/buildspec_vanilla_framework_tests.yml
index 03e1fb317..985990824 100644
--- a/config/buildspec_vanilla_framework_tests.yml
+++ b/config/buildspec_vanilla_framework_tests.yml
@@ -18,7 +18,7 @@ phases:
         - sudo apt-get install unzip -qq -o=Dpkg::Use-Pty=0
         - cd $CODEBUILD_SRC_DIR  && chmod +x config/protoc_downloader.sh && ./config/protoc_downloader.sh
         - pip install --upgrade pip==19.3.1
-        - pip install -q pytest pytest-cov wheel pyYaml pytest-html keras==2.3.1 mxnet torch xgboost pre-commit tensorflow_datasets torchvision
+        - pip install -q pytest pytest-cov wheel pyYaml pytest-html keras==2.3.1 mxnet torch xgboost pre-commit tensorflow_datasets==4.0.1 torchvision
         - cd $CODEBUILD_SRC_DIR && chmod +x config/install_smdebug.sh && chmod +x config/check_smdebug_install.sh && ./config/install_smdebug.sh;
 
   build:
diff --git a/config/install_smdebug.sh b/config/install_smdebug.sh
index d16a09727..cd9e9583a 100644
--- a/config/install_smdebug.sh
+++ b/config/install_smdebug.sh
@@ -73,7 +73,7 @@ if [ "$run_pytest_mxnet" == 'enable' ]; then
 fi
 if [ "$run_pytest_tensorflow" == 'enable' ]; then
   ./config/check_smdebug_install.sh tensorflow
-  pip install tensorflow_datasets
+  pip install tensorflow_datasets==4.0.1
 fi
 if [ "$run_pytest_pytorch" == 'enable' ]; then
   ./config/check_smdebug_install.sh torch
diff --git a/config/tests.sh b/config/tests.sh
index 73e31f4e4..9185840d5 100644
--- a/config/tests.sh
+++ b/config/tests.sh
@@ -66,7 +66,7 @@ if [ "$run_pytest_tensorflow" = "enable" ] ; then
 fi
 
 if [ "$run_pytest_tensorflow2" = "enable" ] ; then
-    pip install tensorflow_datasets
+    pip install tensorflow_datasets==4.0.1
     run_for_framework tensorflow2
 fi