Skip to content

Commit fa2ef2a

Browse files
EddyLXJfacebook-github-bot
authored andcommitted
Fix embedding table type and eviction policy in st publish
Summary: As title, fix st publish gap for eviction Differential Revision: D80729527
1 parent af6f7bf commit fa2ef2a

File tree

1 file changed

+6
-6
lines changed

1 file changed

+6
-6
lines changed

torchrec/distributed/batched_embedding_kernel.py

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -225,12 +225,12 @@ def _populate_zero_collision_tbe_params(
225225
meta_header_lens = [0] * len(config.embedding_tables)
226226
for i, table in enumerate(config.embedding_tables):
227227
# virtual_table_eviction_policy won't be None in reality: https://fburl.com/code/864a0w0f
228-
if table.virtual_table_eviction_policy is not None:
229-
meta_header_lens[i] = (
230-
table.virtual_table_eviction_policy.get_meta_header_len()
231-
)
232-
if not isinstance(table.virtual_table_eviction_policy, NoEvictionPolicy):
233-
enabled = True
228+
assert (
229+
table.virtual_table_eviction_policy is not None
230+
), "virtual_table_eviction_policy for kvzch table should not be None"
231+
meta_header_lens[i] = table.virtual_table_eviction_policy.get_meta_header_len()
232+
if not isinstance(table.virtual_table_eviction_policy, NoEvictionPolicy):
233+
enabled = True
234234
if enabled:
235235
counter_thresholds = [0] * len(config.embedding_tables)
236236
ttls_in_mins = [0] * len(config.embedding_tables)

0 commit comments

Comments
 (0)