From 6330970667f077de08ee379a21c5f4b382430120 Mon Sep 17 00:00:00 2001
From: Jianyu Huang <jianyuhuang@meta.com>
Date: Sat, 18 May 2024 21:06:03 -0700
Subject: [PATCH] resolve genai build fix: still print error info

---
 fbgemm_gpu/fbgemm_gpu/__init__.py      | 14 ++++++++++----
 fbgemm_gpu/fbgemm_gpu/docs/__init__.py |  2 +-
 2 files changed, 11 insertions(+), 5 deletions(-)

diff --git a/fbgemm_gpu/fbgemm_gpu/__init__.py b/fbgemm_gpu/fbgemm_gpu/__init__.py
index 7fffab2ae3..47baa9e722 100644
--- a/fbgemm_gpu/fbgemm_gpu/__init__.py
+++ b/fbgemm_gpu/fbgemm_gpu/__init__.py
@@ -9,11 +9,18 @@
 
 import torch
 
+
 try:
     torch.ops.load_library(os.path.join(os.path.dirname(__file__), "fbgemm_gpu_py.so"))
-except Exception as e:
-    pass
-    # print(e)
+except Exception as error_ranking:
+    try:
+        torch.ops.load_library(
+            os.path.join(os.path.dirname(__file__), "experimental/gen_ai/fbgemm_gpu_experimental_gen_ai_py.so")
+        )
+    except Exception as error_gen_ai:
+        # When both ranking/gen_ai so files are not available, print the error logs
+        print(error_ranking)
+        print(error_gen_ai)
 
 # Since __init__.py is only used in OSS context, we define `open_source` here
 # and use its existence to determine whether or not we are in OSS context
@@ -30,4 +37,3 @@
     from . import sparse_ops  # noqa: F401, E402
 except Exception as e:
     pass
-    # print(e)
diff --git a/fbgemm_gpu/fbgemm_gpu/docs/__init__.py b/fbgemm_gpu/fbgemm_gpu/docs/__init__.py
index 98f6f73dc6..250f9d58e9 100644
--- a/fbgemm_gpu/fbgemm_gpu/docs/__init__.py
+++ b/fbgemm_gpu/fbgemm_gpu/docs/__init__.py
@@ -8,5 +8,5 @@
 # Trigger the manual addition of docstrings to pybind11-generated operators
 try:
     from . import jagged_tensor_ops, table_batched_embedding_ops  # noqa: F401
-except Exception as e:
+except Exception:
     pass