From 6599bae3ae10f52171d4eebbf4bff65d148a56f0 Mon Sep 17 00:00:00 2001
From: Paul Zhang <paulzhan@meta.com>
Date: Thu, 15 Aug 2024 16:48:01 -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

fbshipit-source-id: b565cc01b1c28961cb03b528fd1e951e225b4c24
---
 .../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 8564fa41c6..145f00e96d 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
@@ -47,6 +47,8 @@
 except Exception:
     pass
 
+import fbgemm_gpu  # noqa
+
 
 def rounded_row_size_in_bytes(
     dim: int,