diff --git a/.github/workflows/test_python.yml b/.github/workflows/test_python.yml
index dadd749b36..14cdd4cea0 100644
--- a/.github/workflows/test_python.yml
+++ b/.github/workflows/test_python.yml
@@ -38,16 +38,16 @@ jobs:
             tf: 1.14
           - python: 3.7
             gcc: 5
-            tf: 2.3
+            tf: 2.*
           - python: 3.7
             gcc: 8
-            tf: 2.3
+            tf: 2.*
           - python: 3.8
             gcc: 5
-            tf: 2.3
+            tf: 2.*
           - python: 3.8
             gcc: 8
-            tf: 2.3
+            tf: 2.*
 
     steps:
     - uses: actions/checkout@master
diff --git a/doc/install/install-from-source.md b/doc/install/install-from-source.md
index 9a20a565e3..b0e6f468b1 100644
--- a/doc/install/install-from-source.md
+++ b/doc/install/install-from-source.md
@@ -27,7 +27,7 @@ We follow the virtual environment approach to install the tensorflow's Python in
 virtualenv -p python3 $tensorflow_venv
 source $tensorflow_venv/bin/activate
 pip install --upgrade pip
-pip install --upgrade tensorflow==2.3.0
+pip install --upgrade tensorflow
 ```
 It is notice that everytime a new shell is started and one wants to use `DeePMD-kit`, the virtual environment should be activated by 
 ```bash
@@ -43,7 +43,7 @@ virtualenv -p python3.7 $tensorflow_venv
 ```
 If one does not need the GPU support of deepmd-kit and is concerned about package size, the CPU-only version of tensorflow should be installed by	
 ```bash	
-pip install --upgrade tensorflow-cpu==2.3.0	
+pip install --upgrade tensorflow-cpu
 ```
 To verify the installation, run
 ```bash
@@ -149,4 +149,4 @@ dp_ipi      dp_ipi_low
 $ ls $deepmd_root/lib
 libdeepmd_cc_low.so  libdeepmd_ipi_low.so  libdeepmd_lmp_low.so  libdeepmd_low.so          libdeepmd_op_cuda.so  libdeepmd_op.so
 libdeepmd_cc.so      libdeepmd_ipi.so      libdeepmd_lmp.so      libdeepmd_op_cuda_low.so  libdeepmd_op_low.so   libdeepmd.so
-```
\ No newline at end of file
+```
diff --git a/setup.py b/setup.py
index e48f9a3f3d..618ad59de8 100644
--- a/setup.py
+++ b/setup.py
@@ -21,9 +21,14 @@
 readme_file = Path(__file__).parent / "README.md"
 readme = readme_file.read_text()
 
-tf_version = os.environ.get("TENSORFLOW_VERSION", "2.3")
+tf_version = os.environ.get("TENSORFLOW_VERSION", "")
 
-if tf_version in SpecifierSet("<1.15") or tf_version in SpecifierSet(">=2.0,<2.1"):
+if tf_version == "":
+    extras_require = {
+        "cpu": ["tensorflow-cpu"],
+        "gpu": ["tensorflow"],
+    }
+elif tf_version in SpecifierSet("<1.15") or tf_version in SpecifierSet(">=2.0,<2.1"):
     extras_require = {
         "cpu": [f"tensorflow=={tf_version}"],
         "gpu": [f"tensorflow-gpu=={tf_version}"],