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

fbshipit-source-id: b565cc01b1c28961cb03b528fd1e951e225b4c24

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,