From 55a594e3b8b4eae4fda66a652f05a2d105b08175 Mon Sep 17 00:00:00 2001
From: Julian Qian <julian.qian@gmail.com>
Date: Wed, 28 Aug 2024 13:31:09 -0700
Subject: [PATCH] combine cpu and gpu wheel

---
 .github/workflows/release.yml                 |  2 ++
 README.md                                     |  2 +-
 setup.py                                      |  4 ----
 .../utils/resource_loader.py                  | 21 ++++++++-----------
 4 files changed, 12 insertions(+), 17 deletions(-)

diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml
index 7e46eac10..bac11e74f 100644
--- 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
index 88f289c61..f188383f9 100644
--- 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
index 50475f4d7..1cf6c5a92 100644
--- 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
index a0e5a1c9e..8453311ca 100644
--- 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