diff --git a/.gitignore b/.gitignore
index 83ef8ea38cd3..e24999ef0d5c 100644
--- a/.gitignore
+++ b/.gitignore
@@ -234,6 +234,9 @@ conda/pkg
 .envrc
 *.nix
 
+# Docker files
+.sudo_as_admin_successful
+
 # Downloaded models/datasets
 .tvm_test_data
 .dgl
diff --git a/docker/install/ubuntu_install_redis.sh b/docker/install/ubuntu_install_redis.sh
index 21dab1f3a57b..0eb46eb8edec 100755
--- a/docker/install/ubuntu_install_redis.sh
+++ b/docker/install/ubuntu_install_redis.sh
@@ -21,4 +21,4 @@ set -u
 set -o pipefail
 
 apt-get update && apt-get install -y redis-server
-pip3 install xgboost>=1.1.0 psutil
+pip3 install "xgboost>=1.1.0" psutil
diff --git a/docker/install/ubuntu_install_tensorflow.sh b/docker/install/ubuntu_install_tensorflow.sh
index 286a086abd82..81802964ba0e 100755
--- a/docker/install/ubuntu_install_tensorflow.sh
+++ b/docker/install/ubuntu_install_tensorflow.sh
@@ -20,7 +20,4 @@ set -e
 set -u
 set -o pipefail
 
-# h5py is pinned to minor than 3 due to issues with
-# tensorflow:
-# https://github.com/tensorflow/tensorflow/issues/44467
-pip3 install tensorflow==2.3.1 keras==2.4.3 "h5py<3.0"
+pip3 install tensorflow==2.4.2
diff --git a/docker/install/ubuntu_install_tflite.sh b/docker/install/ubuntu_install_tflite.sh
index 2dfbb0681a80..cfaa643d50f3 100755
--- a/docker/install/ubuntu_install_tflite.sh
+++ b/docker/install/ubuntu_install_tflite.sh
@@ -20,6 +20,10 @@ set -e
 set -u
 set -o pipefail
 
+# The tflite version should have matched versions to the tensorflow
+# version installed from pip in ubuntu_install_tensorflow.sh
+TENSORFLOW_VERSION=$(python3 -c "import tensorflow; print(tensorflow.__version__)" 2> /dev/null)
+
 # Download, build and install flatbuffers
 git clone --branch=v1.12.0 --depth=1 --recursive https://github.com/google/flatbuffers.git
 cd flatbuffers
@@ -33,14 +37,14 @@ pip3 install flatbuffers
 # Build the TFLite static library, necessary for building with TFLite ON.
 # The library is built at:
 # tensorflow/tensorflow/lite/tools/make/gen/*/lib/libtensorflow-lite.a.
-git clone https://github.com/tensorflow/tensorflow --branch=r2.3
+git clone https://github.com/tensorflow/tensorflow --branch=v${TENSORFLOW_VERSION}
 ./tensorflow/tensorflow/lite/tools/make/download_dependencies.sh
 ./tensorflow/tensorflow/lite/tools/make/build_lib.sh
 
 # Setup tflite from schema
 mkdir tflite
+cp tensorflow/tensorflow/lite/schema/schema.fbs tflite
 cd tflite
-wget -q https://raw.githubusercontent.com/tensorflow/tensorflow/r2.3/tensorflow/lite/schema/schema.fbs
 flatc --python schema.fbs
 
 cat <<EOM >setup.py
@@ -48,7 +52,7 @@ import setuptools
 
 setuptools.setup(
     name="tflite",
-    version="2.3.1",
+    version="${TENSORFLOW_VERSION}",
     author="google",
     author_email="google@google.com",
     description="TFLite",