diff --git a/config/buildspec.yml b/config/buildspec.yml
index fefc3579d..17a30113b 100755
--- a/config/buildspec.yml
+++ b/config/buildspec.yml
@@ -30,7 +30,7 @@ phases:
         - cd $CODEBUILD_SRC_DIR  && chmod +x config/protoc_downloader.sh && ./config/protoc_downloader.sh
         - pip install --upgrade pip==19.3.1
         - if [ "$run_pytest_xgboost" = "enable" ]; then pip install --upgrade pyYaml==5.1; else pip install -q pyYaml; fi
-        - pip install -q pytest==5.3.3 wheel pytest-html pre-commit awscli tensorflow_datasets pytest-cov
+        - pip install -q pytest wheel pytest-html pre-commit awscli pytest-cov
 
   pre_build:
     commands:
@@ -39,7 +39,8 @@ phases:
   build:
     commands:
       - cd $CODEBUILD_SRC_DIR  && python setup.py bdist_wheel --universal
-      - cd $RULES_CODEBUILD_SRC_DIR && python setup.py bdist_wheel --universal && pip install --force-reinstall dist/*.whl && cd ..
+      - cd $RULES_CODEBUILD_SRC_DIR && python setup.py bdist_wheel --universal
+      - if [ "$run_pytest_xgboost" = "enable" ]; then pip install --force-reinstall $RULES_CODEBUILD_SRC_DIR/dist/*.whl; else pip install $RULES_CODEBUILD_SRC_DIR/dist/*.whl; fi
       - cd $CODEBUILD_SRC_DIR  && pip install --force-reinstall dist/*.whl && cd ..
       - cd $CODEBUILD_SRC_DIR  && chmod +x config/tests.sh && PYTHONPATH=. && ./config/tests.sh  && mkdir -p upload/$CURRENT_COMMIT_PATH/wheels && cp ./dist/*.whl upload/$CURRENT_COMMIT_PATH/wheels && cd ..
       - pip show smdebug
diff --git a/config/buildspec_build_wheel.yml b/config/buildspec_build_wheel.yml
index 6d7320ae9..421ce8619 100644
--- a/config/buildspec_build_wheel.yml
+++ b/config/buildspec_build_wheel.yml
@@ -16,7 +16,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==5.3.3 wheel pyYaml pytest-html pre-commit pytest-cov
+        - pip install -q pytest wheel pyYaml pytest-html pre-commit pytest-cov
         - pip uninstall -y boto3 && pip uninstall -y aiobotocore && pip uninstall -y botocore
 
   build:
diff --git a/config/buildspec_vanilla_framework_tests.yml b/config/buildspec_vanilla_framework_tests.yml
index 09820ed51..03e1fb317 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==5.3.3 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 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/buildspec_zero_code_change.yml b/config/buildspec_zero_code_change.yml
index f8a5db709..9958ec3d3 100644
--- a/config/buildspec_zero_code_change.yml
+++ b/config/buildspec_zero_code_change.yml
@@ -25,7 +25,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==5.3.3 wheel pyYaml pytest-html pre-commit awscli tensorflow_datasets pytest-cov
+      - pip install -q pytest wheel pyYaml pytest-html pre-commit awscli pytest-cov
       - 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/tests.sh b/config/tests.sh
index d08aec9d6..73e31f4e4 100644
--- a/config/tests.sh
+++ b/config/tests.sh
@@ -60,10 +60,13 @@ if [ "$run_pytest_xgboost" = "enable" ] ; then
 fi
 
 if [ "$run_pytest_tensorflow" = "enable" ] ; then
+    # tensorflow_datasets >= 4.0.0 requires tensorflow >= 2.1.0
+    pip install tensorflow_datasets==3.2.1
     run_for_framework tensorflow
 fi
 
 if [ "$run_pytest_tensorflow2" = "enable" ] ; then
+    pip install tensorflow_datasets
     run_for_framework tensorflow2
 fi