2025-01-14 nightly release (542b0b2)

pytorch · Jan 14, 2025 · 42fa9d0 · 42fa9d0
1 parent 6f54403
commit 42fa9d0
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 0 deletions.
diff --git a/torchrec/distributed/batched_embedding_kernel.py b/torchrec/distributed/batched_embedding_kernel.py
@@ -624,6 +624,9 @@ def step(self, closure: Any = None) -> None:
     def set_optimizer_step(self, step: int) -> None:
         self._emb_module.set_optimizer_step(step)
 
+    def update_hyper_parameters(self, params_dict: Dict[str, Any]) -> None:
+        self._emb_module.update_hyper_parameters(params_dict)
+
 
 def _gen_named_parameters_by_table_ssd(
     emb_module: SSDTableBatchedEmbeddingBags,

diff --git a/torchrec/distributed/global_settings.py b/torchrec/distributed/global_settings.py
@@ -30,3 +30,7 @@ def construct_sharded_tensor_from_metadata_enabled() -> bool:
     return (
         os.environ.get(TORCHREC_CONSTRUCT_SHARDED_TENSOR_FROM_METADATA_ENV, "0") == "1"
     )
+
+
+def enable_construct_sharded_tensor_from_metadata() -> None:
+    os.environ[TORCHREC_CONSTRUCT_SHARDED_TENSOR_FROM_METADATA_ENV] = "1"