2024-11-23 nightly release (7f3b7dc)

pytorch · Nov 23, 2024 · c5072a2 · c5072a2
1 parent 5560df4
commit c5072a2
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 4 deletions.
diff --git a/torchrec/distributed/embeddingbag.py b/torchrec/distributed/embeddingbag.py
@@ -953,9 +953,9 @@ def _initialize_torch_state(self) -> None:  # noqa
                 )
 
                 self._model_parallel_name_to_sharded_tensor[table_name] = (
-                    ShardedTensor._init_from_local_shards(
-                        local_shards,
-                        self._name_to_table_size[table_name],
+                    ShardedTensor._init_from_local_shards_and_global_metadata(
+                        local_shards=local_shards,
+                        sharded_tensor_metadata=metadata,
                         process_group=(
                             self._env.sharding_pg
                             if isinstance(self._env, ShardingEnv2D)

diff --git a/torchrec/pt2/utils.py b/torchrec/pt2/utils.py
@@ -75,12 +75,15 @@ def kjt_for_pt2_tracing(
 
     values = kjt.values().long()
     torch._dynamo.decorators.mark_unbacked(values, 0)
+    weights = kjt.weights_or_none()
+    if weights is not None:
+        torch._dynamo.decorators.mark_unbacked(weights, 0)
 
     return KeyedJaggedTensor(
         keys=kjt.keys(),
         values=values,
         lengths=lengths,
-        weights=kjt.weights_or_none(),
+        weights=weights,
         stride=stride if not is_vb else None,
         stride_per_key_per_rank=kjt.stride_per_key_per_rank() if is_vb else None,
         inverse_indices=inverse_indices,