Skip to content

Commit

Permalink
pass table names to TBE for EC as well
Browse files Browse the repository at this point in the history
Summary: followup to D54442088, which passed table names to EBC but not EC (blame is on me, totally forgot about that)

Differential Revision: D56076712
  • Loading branch information
henrylhtsang authored and facebook-github-bot committed Apr 12, 2024
1 parent 568e116 commit 4d7431f
Showing 1 changed file with 1 addition and 0 deletions.
1 change: 1 addition & 0 deletions torchrec/distributed/batched_embedding_kernel.py
Original file line number Diff line number Diff line change
Expand Up @@ -604,6 +604,7 @@ def __init__(
pooling_mode=PoolingMode.NONE,
weights_precision=weights_precision,
device=device,
table_names=[t.name for t in config.embedding_tables],
**fused_params,
)
)
Expand Down

0 comments on commit 4d7431f

Please sign in to comment.