diff --git a/WORKSPACE b/WORKSPACE
index ff69268a4..b88b0d0b0 100644
--- a/WORKSPACE
+++ b/WORKSPACE
@@ -1,6 +1,6 @@
 # This file includes external dependencies that are required to compile the
 # TensorFlow op. Maybe of them are specific versions used by the TensorFlow
-# binary used. These are extracted from TF v2.3.1, tensorflow/workspace.bzl.
+# binary used. These are extracted from TF v2.4.1, tensorflow/workspace.bzl.
 
 load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive")
 
@@ -85,10 +85,10 @@ http_archive(
 
 http_archive(
     name = "org_tensorflow",
-    sha256 = "6f063636673d6ef4ac60febd2541e3ad3516a57c18339a680c794b736798d054",
-    strip_prefix = "tensorflow-2.3.1",
+    sha256 = "b91ec194ddf6c4a5a2f9d1db4af4daab0b187ff691e6f88142413d2c7e77a3bb",
+    strip_prefix = "tensorflow-2.4.1",
     urls = [
-        "https://github.com/tensorflow/tensorflow/archive/v2.3.1.zip",
+        "https://github.com/tensorflow/tensorflow/archive/v2.4.1.zip",
     ],
 )
 
diff --git a/configure.sh b/configure.sh
index 4923c983a..2af283951 100755
--- a/configure.sh
+++ b/configure.sh
@@ -72,47 +72,25 @@ while [[ "$TF_CUDA_VERSION" == "" ]]; do
 done
 
 
-# CPU
-if [[ "$TF_NEED_CUDA" == "0" ]]; then
-
-  # Check if it's installed
-  if [[ $(pip show tensorflow-cpu) == *tensorflow-cpu* ]] || [[ $(pip show tf-nightly-cpu) == *tf-nightly-cpu* ]] ; then
-    echo 'Using installed tensorflow'
-  else
-    # Uninstall GPU version if it is installed.
-    if [[ $(pip show tensorflow) == *tensorflow* ]]; then
-      echo 'Already have gpu version of tensorflow installed. Uninstalling......\n'
-      pip uninstall tensorflow
-    elif [[ $(pip show tf-nightly) == *tf-nightly* ]]; then
-      echo 'Already have gpu version of tensorflow installed. Uninstalling......\n'
-      pip uninstall tf-nightly
-    fi
-    # Install CPU version
-    echo 'Installing tensorflow-cpu......\n'
-    pip install tensorflow-cpu
-  fi
-
+# Check if it's installed
+if [[ $(pip show tensorflow) == *tensorflow* ]] || [[ $(pip show tf-nightly) == *tf-nightly* ]]; then
+  echo 'Using installed tensorflow'
 else
-
-  # Check if it's installed
-   if [[ $(pip show tensorflow) == *tensorflow* ]] || [[ $(pip show tf-nightly) == *tf-nightly* ]]; then
-    echo 'Using installed tensorflow'
-  else
-    # Uninstall CPU version if it is installed.
-    if [[ $(pip show tensorflow-cpu) == *tensorflow-cpu* ]]; then
-      echo 'Already have tensorflow non-gpu installed. Uninstalling......\n'
-      pip uninstall tensorflow
-    elif [[ $(pip show tf-nightly-cpu) == *tf-nightly-cpu* ]]; then
-      echo 'Already have tensorflow non-gpu installed. Uninstalling......\n'
-      pip uninstall tf-nightly
-    fi
-    # Install GPU version
-    echo 'Installing tensorflow .....\n'
-    pip install tensorflow
+  # Uninstall CPU version if it is installed.
+  if [[ $(pip show tensorflow-cpu) == *tensorflow-cpu* ]]; then
+    echo 'Already have tensorflow non-gpu installed. Uninstalling......\n'
+    pip uninstall tensorflow
+  elif [[ $(pip show tf-nightly-cpu) == *tf-nightly-cpu* ]]; then
+    echo 'Already have tensorflow non-gpu installed. Uninstalling......\n'
+    pip uninstall tf-nightly
   fi
+  # Install GPU version
+  echo 'Installing tensorflow .....\n'
+  pip install tensorflow
 fi
 
 
+
 TF_CFLAGS=( $(python -c 'import tensorflow as tf; print(" ".join(tf.sysconfig.get_compile_flags()))') )
 TF_LFLAGS="$(python -c 'import tensorflow as tf; print(" ".join(tf.sysconfig.get_link_flags()))')"
 
diff --git a/release/setup.py b/release/setup.py
index 155d835b3..8138169fb 100644
--- a/release/setup.py
+++ b/release/setup.py
@@ -54,7 +54,7 @@ def finalize_options(self):
 
 # placed as extra to not have required overwrite existing nightly installs if
 # they exist.
-EXTRA_PACKAGES = ['tensorflow == 2.3.1']
+EXTRA_PACKAGES = ['tensorflow == 2.4.1']
 CUR_VERSION = '0.5.0'
 
 
diff --git a/requirements.txt b/requirements.txt
index 28deb4362..2d21bfb32 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -4,5 +4,5 @@ nbconvert==5.6.1
 nbformat==4.4.0
 pylint==2.4.4
 yapf==0.28.0
-tensorflow==2.3.1
+tensorflow==2.4.1
 google-api-python-client==1.8.0