Skip to content

Commit a5dba57

Browse files
EddyLXJmeta-codesync[bot]
authored andcommitted
Fix is_using_virtual_table bug in kernel
Summary: Ease fix for a misused func calling. `config.is_using_virtual_table` will always be true. Reviewed By: emlin Differential Revision: D86233647 fbshipit-source-id: 6d91b5f104c7cc098998d012245823991372bc87
1 parent 2996a9f commit a5dba57

File tree

1 file changed

+4
-4
lines changed

1 file changed

+4
-4
lines changed

torchrec/distributed/batched_embedding_kernel.py

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -194,7 +194,7 @@ def _populate_ssd_tbe_params(config: GroupedEmbeddingConfig) -> Dict[str, Any]:
194194
)
195195

196196
# populate init min and max
197-
if config.is_using_virtual_table:
197+
if config.is_using_virtual_table():
198198
_generate_init_range_for_virtual_tables(ssd_tbe_params, config)
199199

200200
if (
@@ -242,7 +242,7 @@ def _populate_ssd_tbe_params(config: GroupedEmbeddingConfig) -> Dict[str, Any]:
242242
)
243243
ssd_tbe_params["cache_sets"] = int(max_cache_sets)
244244

245-
if "kvzch_eviction_tbe_config" in fused_params and config.is_using_virtual_table:
245+
if "kvzch_eviction_tbe_config" in fused_params and config.is_using_virtual_table():
246246
ssd_tbe_params["kvzch_eviction_tbe_config"] = fused_params.get(
247247
"kvzch_eviction_tbe_config"
248248
)
@@ -1969,7 +1969,7 @@ def __init__(
19691969
len({table.embedding_dim for table in config.embedding_tables}) == 1
19701970
), "Currently we expect all tables in SSD TBE to have the same embedding dimension."
19711971
assert (
1972-
config.is_using_virtual_table
1972+
config.is_using_virtual_table()
19731973
), "Try to create ZeroCollisionKeyValueEmbedding for non virtual tables"
19741974
assert embedding_cache_mode == config.enable_embedding_update, (
19751975
f"Embedding_cache kernel is {embedding_cache_mode} "
@@ -2883,7 +2883,7 @@ def __init__(
28832883
len({table.embedding_dim for table in config.embedding_tables}) == 1
28842884
), "Currently we expect all tables in SSD TBE to have the same embedding dimension."
28852885
assert (
2886-
config.is_using_virtual_table
2886+
config.is_using_virtual_table()
28872887
), "Try to create ZeroCollisionKeyValueEmbeddingBag for non virtual tables"
28882888

28892889
for table in config.embedding_tables:

0 commit comments

Comments
 (0)