Fix metadata tensor size and add ut to cover that. (#4704)

EddyLXJ · facebook-github-bot · commit 2ac0b217d3e4 · 2025-08-14T16:09:13.000-07:00
Summary: Pull Request resolved: #4704 X-link: facebookresearch/FBGEMM#1729 As title, currently has size mismatch issue {F1981225318} Reviewed By: kathyxuyy Differential Revision: D80282519 fbshipit-source-id: df9e7a588bb7d73aaa0a993c98d669aaebc7761d
diff --git a/fbgemm_gpu/fbgemm_gpu/tbe/ssd/training.py b/fbgemm_gpu/fbgemm_gpu/tbe/ssd/training.py
@@ -2982,7 +2982,7 @@ def split_embedding_weights(
                     bucket_ascending_id_tensor + table_offset,
                     torch.as_tensor(bucket_ascending_id_tensor.size(0)),
                     snapshot_handle,
-                )
+                ).view(-1, 1)
 
                 # 3. convert local id back to global id
                 bucket_ascending_id_tensor.add_(bucket_id_start * bucket_size)
diff --git a/fbgemm_gpu/test/tbe/ssd/ssd_split_tbe_training_test.py b/fbgemm_gpu/test/tbe/ssd/ssd_split_tbe_training_test.py
@@ -2250,6 +2250,8 @@ def test_kv_emb_state_dict(
                 rtol=tolerance,
             )
 
+            self.assertTrue(len(metadata_list[table_index].size()) == 2)
+
     @given(
         **{
             "T": st.integers(min_value=1, max_value=10),

Original file line number	Diff line number	Diff line change
`@@ -2250,6 +2250,8 @@ def test_kv_emb_state_dict(`
`2250`	`2250`	`rtol=tolerance,`
`2251`	`2251`	`)`
`2252`	`2252`
	`2253`	`+ self.assertTrue(len(metadata_list[table_index].size()) == 2)`
	`2254`	`+`
`2253`	`2255`	`@given(`
`2254`	`2256`	`**{`
`2255`	`2257`	`"T": st.integers(min_value=1, max_value=10),`