Add logging for VBE usage in torchrec (#1486)

Summary: Pull Request resolved: #1486 Add logging for vbe usage in torchrec. Reviewed By: colin2328 Differential Revision: D51049595 fbshipit-source-id: 3fc5e280b40f3afc3da39f5250e69e742cfef317
pytorch · Nov 7, 2023 · e3e9661 · e3e9661
1 parent 915dfc5
commit e3e9661
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/torchrec/distributed/dist_data.py b/torchrec/distributed/dist_data.py
@@ -733,6 +733,7 @@ def __init__(
         codecs: Optional[QuantizedCommCodecs] = None,
     ) -> None:
         super().__init__()
+        torch._C._log_api_usage_once("torchrec.distributed.vbe")
         self._pg = pg
         self._emb_dim_per_rank_per_feature = emb_dim_per_rank_per_feature
         self._callbacks: List[Callable[[torch.Tensor], torch.Tensor]] = []