Improve Inference UX, add tests, and add inference API docstrings (#2988

) Summary: X-link: facebookresearch/FBGEMM#81 Pull Request resolved: #2988 X-link: pytorch/torchrec#2295 Reviewed By: jingsh Differential Revision: D61175593
pytorch · Aug 15, 2024 · e990925 · e990925
1 parent 7105d5f
commit e990925
Showing 1 changed file with 2 additions and 0 deletions.
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
@@ -46,6 +46,8 @@
 except Exception:
     pass
 
+import fbgemm_gpu  # noqa
+
 
 def rounded_row_size_in_bytes(
     dim: int,