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,