setup

Jenkinsfile

@@ -45,6 +45,5 @@ pipeline {
         }
       }
     }
-
   }
 }

python/setup.py

@@ -85,7 +85,6 @@ def build_extension(self, ext):
     with open("MANIFEST.in", "w") as fo:
         for path in LIBS:
             shutil.copy(path, os.path.join(CURRENT_DIR, 'tfdlpack'))
-            print('>>>>>>>>>>>>>', path)
             _, libname = os.path.split(path)
             fo.write("include tfdlpack/%s\n" % libname)
     setup_kwargs = {
@@ -108,7 +107,7 @@ def build_extension(self, ext):
     author_email='allen.zhou@nyu.edu',
     description='Tensorflow plugin for DLPack',
     packages=find_packages(),
-    install_requires=['tensorflow>=2.0.0'],
+    install_requires=['tensorflow%s>=2.0.0' % os.getenv('TFDLPACK_PACKAGE_SUFFIX')],
     long_description="""
 The package adds interoperability of DLPack to Tensorflow. It contains straightforward
 and easy-to-use APIs to convert Tensorflow tensors from/to DLPack format.

tests/scripts/task_build.sh

@@ -14,6 +14,7 @@ make -j4
 popd
 
 export TFDLPACK_LIBRARY_PATH=$PWD/build
+export TFDLPACK_PACKAGE_SUFFIX=-gpu
 
 pushd python
 python3 setup.py clean