combine cpu and gpu wheel

tensorflow · Aug 29, 2024 · 55a594e · 55a594e
1 parent 0dd0fd8
commit 55a594e
Show file tree

Hide file tree

Showing 4 changed files with 12 additions and 17 deletions.
diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml
@@ -207,6 +207,8 @@ jobs:
             py-version: '3.8'
           - os: 'Linux'
             cpu: 'arm64'
+          - os: 'Linux'
+            tf-need-cuda: '0'
           - py-version: '3.7'
             cpu: 'arm64'
           - py-version: '3.8'

diff --git a/README.md b/README.md
@@ -76,7 +76,7 @@ run the following:
 pip install tensorflow-recommenders-addons
 ```
 
-By default, CPU version will be installed. To install GPU version, run the following:
+Before the release of combined CPU and GPU version, to install GPU version, run the following:
 ```
 pip install tensorflow-recommenders-addons-gpu
 ```

diff --git a/setup.py b/setup.py
@@ -55,13 +55,9 @@ def get_project_name_version():
 
   project_name = "tensorflow-recommenders-addons"
   version["tf_project_name"] = "tensorflow"
-  if os.getenv("TF_NEED_CUDA", "0") == "1":
-    project_name = project_name + "-gpu"
 
   if "--nightly" in sys.argv:
     project_name = "tfra-nightly"
-    if os.getenv("TF_NEED_CUDA", "0") == "1":
-      project_name = project_name + "-gpu"
     version["__version__"] += get_last_commit_time()
     sys.argv.remove("--nightly")
 

diff --git a/tensorflow_recommenders_addons/utils/resource_loader.py b/tensorflow_recommenders_addons/utils/resource_loader.py
@@ -27,18 +27,15 @@
 def get_required_tf_version():
   try:
     pkg = pkg_resources.get_distribution("tensorflow-recommenders-addons")
-  except:
-    try:
-      pkg = pkg_resources.get_distribution("tensorflow-recommenders-addons-gpu")
-    except pkg_resources.DistributionNotFound:
-      # Force return for 'Test with bazel' on CI.
-      warnings.warn(
-          "Fail to get TFRA package information, if you are running on "
-          "bazel test mode, please ignore this warning, \nor you should check "
-          "TFRA installation.",
-          UserWarning,
-      )
-      return tf.__version__, tf.__version__
+  except pkg_resources.DistributionNotFound:
+    # Force return for 'Test with bazel' on CI.
+    warnings.warn(
+        "Fail to get TFRA package information, if you are running on "
+        "bazel test mode, please ignore this warning, \nor you should check "
+        "TFRA installation.",
+        UserWarning,
+    )
+    return tf.__version__, tf.__version__
 
   pkg_info = pkg.requires()
   low_version, high_version = None, None