Fix test for embedding sharding (#4456)

kausv · facebook-github-bot · commit b4137fc4be11 · 2025-07-09T08:28:47.000-07:00
Summary: X-link: facebookresearch/FBGEMM#1517 Pull Request resolved: #4456 We were trying to log an optional tensor. I added conditional logging. Failing test https://www.internalfb.com/intern/test/281475196311912 Reviewed By: q10 Differential Revision: D77963539 fbshipit-source-id: a79f7cae51a6b38ee2dee25b6e1be2e571865423
diff --git a/fbgemm_gpu/fbgemm_gpu/tbe/ssd/training.py b/fbgemm_gpu/fbgemm_gpu/tbe/ssd/training.py
@@ -2291,8 +2291,7 @@ def split_optimizer_states(
             table_offset += emb_height
         logging.info(
             f"KV ZCH tables split_optimizer_states query latency: {(time.time() - start_time) * 1000} ms, "
-            # pyre-ignore [16]
-            f"num ids list: {[ids.numel() for ids in sorted_id_tensor]}"
+            f"num ids list: {None if not sorted_id_tensor else [ids.numel() for ids in sorted_id_tensor]}"
         )
         return opt_list
 

Original file line number	Diff line number	Diff line change
`@@ -2291,8 +2291,7 @@ def split_optimizer_states(`
`2291`	`2291`	`table_offset += emb_height`
`2292`	`2292`	`logging.info(`
`2293`	`2293`	`f"KV ZCH tables split_optimizer_states query latency: {(time.time() - start_time) * 1000} ms, "`
`2294`		`- # pyre-ignore [16]`
`2295`		`- f"num ids list: {[ids.numel() for ids in sorted_id_tensor]}"`
	`2294`	`+ f"num ids list: {None if not sorted_id_tensor else [ids.numel() for ids in sorted_id_tensor]}"`
`2296`	`2295`	`)`
`2297`	`2296`	`return opt_list`
`2298`	`2297`