More detailed logging in merge_pooled_embedding

Summary: As titled Differential Revision: D58014730
pytorch · May 31, 2024 · b4dc1a7 · b4dc1a7
1 parent cdccc7b
commit b4dc1a7
Showing 1 changed file with 11 additions and 1 deletion.
diff --git a/fbgemm_gpu/src/merge_pooled_embedding_ops/merge_pooled_embedding_ops_gpu.cpp b/fbgemm_gpu/src/merge_pooled_embedding_ops/merge_pooled_embedding_ops_gpu.cpp
@@ -532,7 +532,17 @@ cat_dim_2d_output_shape(
   for (const auto& t : tensors) {
     TORCH_CHECK(t.dim() == 2);
     // only support two-dimension tensors.
-    TORCH_CHECK(t.size(1 - cat_dim) == uncat_dim_size);
+    TORCH_CHECK(
+        t.size(1 - cat_dim) == uncat_dim_size,
+        "cat dim does not match, ",
+        "compare tensor dim = ",
+        1 - cat_dim,
+        ", expect size = ",
+        uncat_dim_size,
+        ", but got size = ",
+        t.size(1 - cat_dim),
+        " on device ",
+        t.get_device());
     total_cat_dim += t.size(cat_dim);
     cumulative_dims.push_back(total_cat_dim);
   }