diff --git a/config/install_smdebug.sh b/config/install_smdebug.sh
index 2f1760d1b..61f951523 100644
--- a/config/install_smdebug.sh
+++ b/config/install_smdebug.sh
@@ -8,7 +8,7 @@ cd $CODEBUILD_SRC_DIR
 if [ "$SMDEBUG_S3_BINARY" ]; then
   mkdir -p s3_pip_binary
   aws s3 cp "$SMDEBUG_S3_BINARY" s3_pip_binary
-  pip install --force-reinstall --upgrade s3_pip_binary/*.whl
+  pip install --upgrade s3_pip_binary/*.whl
 else
   python setup.py bdist_wheel --universal && pip install --upgrade --force-reinstall dist/*.whl
 fi
@@ -19,7 +19,12 @@ if [ "$run_pytest_mxnet" == 'enable' ]; then
 fi
 if [ "$run_pytest_tensorflow" == 'enable' ]; then
   ./config/check_smdebug_install.sh tensorflow
+  pip install tensorflow_datasets
 fi
 if [ "$run_pytest_pytorch" == 'enable' ]; then
   ./config/check_smdebug_install.sh torch
 fi
+if [ "$run_pytest_xgboost" == 'enable' ]; then
+  ./config/check_smdebug_install.sh xgboost
+  pip install --ignore-installed PyYAML
+fi