From e990925b1ff0aad325a5520721a540b094f7e7a4 Mon Sep 17 00:00:00 2001
From: Paul Zhang <paulzhan@meta.com>
Date: Thu, 15 Aug 2024 13:03:05 -0700
Subject: [PATCH] Improve Inference UX, add tests, and add inference API
 docstrings (#2988)

Summary:
X-link: https://github.com/facebookresearch/FBGEMM/pull/81

Pull Request resolved: https://github.com/pytorch/FBGEMM/pull/2988

X-link: https://github.com/pytorch/torchrec/pull/2295

Reviewed By: jingsh

Differential Revision: D61175593
---
 .../fbgemm_gpu/split_table_batched_embeddings_ops_inference.py  | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/fbgemm_gpu/fbgemm_gpu/split_table_batched_embeddings_ops_inference.py b/fbgemm_gpu/fbgemm_gpu/split_table_batched_embeddings_ops_inference.py
index f99ad00530..e7b94c456e 100644
--- a/fbgemm_gpu/fbgemm_gpu/split_table_batched_embeddings_ops_inference.py
+++ b/fbgemm_gpu/fbgemm_gpu/split_table_batched_embeddings_ops_inference.py
@@ -46,6 +46,8 @@
 except Exception:
     pass
 
+import fbgemm_gpu  # noqa
+
 
 def rounded_row_size_in_bytes(
     dim: int,