diff --git a/.github/workflows/test_python.yml b/.github/workflows/test_python.yml
index 14cdd4cea0..8517560fc9 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.*
+            tf:
           - python: 3.7
             gcc: 8
-            tf: 2.*
+            tf:
           - python: 3.8
             gcc: 5
-            tf: 2.*
+            tf:
           - python: 3.8
             gcc: 8
-            tf: 2.*
+            tf:
 
     steps:
     - uses: actions/checkout@master
@@ -70,4 +70,5 @@ jobs:
         CC: gcc-${{ matrix.gcc }}
         CXX: g++-${{ matrix.gcc }}
         TENSORFLOW_VERSION: ${{ matrix.tf }}
+    - run: dp --version
     - run: pytest --cov=deepmd source/tests && codecov
diff --git a/setup.py b/setup.py
index 54a9b36832..e08f587e45 100644
--- a/setup.py
+++ b/setup.py
@@ -8,7 +8,6 @@
 from sysconfig import get_path
 
 from packaging.specifiers import SpecifierSet
-from pkg_resources import Distribution
 from skbuild import setup
 from skbuild.cmaker import get_cmake_version
 from skbuild.exceptions import SKBuildError
@@ -72,11 +71,9 @@
     # TypeError if submodule_search_locations are None
     # IndexError if submodule_search_locations is an empty list
 except (AttributeError, TypeError, IndexError):
-    setup_requires.append(f"tensorflow=={tf_version}")
-    dist = Distribution(
-        project_name="tensorflow", version=tf_version, platform=get_platform()
-    ).egg_name()
-    tf_install_dir = Path(__file__).parent.resolve().joinpath(".egg", dist, "tensorflow").resolve()
+    setup_requires.extend(extras_require['cpu'])
+    # setuptools will re-find tensorflow after installing setup_requires
+    tf_install_dir = None
 
 # add cmake as a build requirement if cmake>3.7 is not installed
 try: