diff --git a/.travis.yml b/.travis.yml
index dd4b0eed9e..b434af6819 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -66,34 +66,30 @@ jobs:
       env:
       - CC=gcc-5
       - CXX=g++-5
-      - TENSORFLOW_VERSION=2.1
+      - TENSORFLOW_VERSION=2.3
     - python: 3.7
       env:
       - CC=gcc-8
       - CXX=g++-8
-      - TENSORFLOW_VERSION=2.1
+      - TENSORFLOW_VERSION=2.3
     - stage: build whls
       services: docker
       env:
         - TWINE_USERNAME=__token__
         - CIBW_BUILD="cp36-* cp37-*"
-        - CIBW_BEFORE_BUILD="pip install tensorflow && sed -i 's/libresolv.so.2\"/libresolv.so.2\", \"libtensorflow_framework.so.2\"/g' \$(find / -name policy.json)"
+        - CIBW_BEFORE_BUILD="sed -i 's/libresolv.so.2\"/libresolv.so.2\", \"libtensorflow_framework.so.2\"/g' \$(find / -name policy.json)"
         - CIBW_SKIP="*-win32 *-manylinux_i686"
         - CC=gcc-7
         - CXX=g++-7
-        - TENSORFLOW_VERSION=2.1
+        - TENSORFLOW_VERSION=2.3
       install:
-        - python -m pip install twine cibuildwheel==1.1.0 scikit-build setuptools_scm
+        - python -m pip install twine cibuildwheel==1.6.3 scikit-build setuptools_scm
       script:
         - python -m cibuildwheel --output-dir wheelhouse
         - python setup.py sdist
       after_success:
         - if [[ $TRAVIS_TAG ]]; then python -m twine upload wheelhouse/*; python -m twine upload dist/*.tar.gz;  fi
-before_install:
-  #- pip install --upgrade pip
-  - pip install --upgrade setuptools
-  - pip install tensorflow==$TENSORFLOW_VERSION
 install:
-  - pip install --verbose .[test]
+  - pip install .[cpu,test]
 script:
   - cd source/tests && python -m unittest
diff --git a/README.md b/README.md
index 5862ed40b5..0d33c8bd12 100644
--- a/README.md
+++ b/README.md
@@ -133,7 +133,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.1.0
+pip install --upgrade tensorflow==2.3.0
 ```
 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
@@ -149,7 +149,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.1.0	
+pip install --upgrade tensorflow-cpu==2.3.0	
 ```
 To verify the installation, run
 ```bash
diff --git a/setup.py b/setup.py
index 117ccec2a2..db86971589 100644
--- a/setup.py
+++ b/setup.py
@@ -4,7 +4,9 @@
 from setuptools_scm import get_version
 from packaging.version import LegacyVersion
 from os import path, makedirs
-import os, imp, sys, platform, sysconfig
+import os, importlib
+import pkg_resources
+from distutils.util import get_platform
 
 
 readme_file = path.join(path.dirname(path.abspath(__file__)), 'README.md')
@@ -15,19 +17,32 @@
     with open(readme_file) as f:
         readme = f.read()
 
-try:
-    tf_install_dir = imp.find_module('tensorflow')[1] 
-except ImportError:
-    site_packages_path = path.join(path.dirname(path.__file__), 'site-packages')
-    tf_install_dir = imp.find_module('tensorflow', [site_packages_path])[1]
-
-
 install_requires=['numpy', 'scipy', 'pyyaml']
-setup_requires=['setuptools_scm', 'scikit-build', 'cmake']
+setup_requires=['setuptools_scm', 'scikit-build']
+
+tf_version = os.environ.get('TENSORFLOW_VERSION', '2.3')
+if LegacyVersion(tf_version) < LegacyVersion("1.15") or (LegacyVersion(tf_version) >= LegacyVersion("2.0") and LegacyVersion(tf_version) <  LegacyVersion("2.1")):
+    extras_require = {"cpu": ["tensorflow==" + tf_version], "gpu": ["tensorflow-gpu==" + tf_version]}
+else:
+    extras_require = {"cpu": ["tensorflow-cpu==" + tf_version], "gpu": ["tensorflow==" + tf_version]}
+tf_spec = importlib.util.find_spec("tensorflow")
+if tf_spec:
+    tf_install_dir = tf_spec.submodule_search_locations[0]
+else:
+    site_packages_path = path.join(path.dirname(path.__file__), 'site-packages')
+    tf_spec = importlib.machinery.FileFinder(site_packages_path).find_spec("tensorflow")
+    if tf_spec:
+        tf_install_dir = tf_spec.submodule_search_locations[0]
+    else:
+        setup_requires.append("tensorflow==" + tf_version)
+        tf_install_dir = path.join(path.dirname(path.abspath(__file__)), '.egg',
+                                   pkg_resources.Distribution(project_name="tensorflow", version=tf_version,
+                                                              platform=get_platform()).egg_name(),
+                                   'tensorflow')
 
-# add cmake as a build requirement if cmake>3.0 is not installed
+# add cmake as a build requirement if cmake>3.7 is not installed
 try:
-    if LegacyVersion(get_cmake_version()) < LegacyVersion("3.0"):
+    if LegacyVersion(get_cmake_version()) < LegacyVersion("3.7"):
         setup_requires.append('cmake')
 except SKBuildError:
     setup_requires.append('cmake')
@@ -65,6 +80,7 @@
     cmake_minimum_required_version='3.0',
     extras_require={
         'test': ['dpdata>=0.1.9'],
+        **extras_require,
     },
     entry_points={
           'console_scripts': ['dp = deepmd.main:main']