Relax the checks for dimensions of pooled_embs (#4165)

Summary:
Pull Request resolved: #4165

- Relax the checks for dimensions of pooled_embs

Reviewed By: spcyppt

Differential Revision: D75158424

fbshipit-source-id: 175bf88283d0a566bba7a97ada36194602e44b8e

Author: q10

fbgemm_gpu/src/permute_pooled_embedding_ops/permute_pooled_embedding_ops.cu

@@ -70,8 +70,8 @@ Tensor permute_pooled_embs_gpu_impl(
   }
 
   TORCH_CHECK(
-      pooled_embs.dim() == 2,
-      "pooled_embs must be 2-D tensor of size [B_local][Sum_T_global(D)], "
+      pooled_embs.dim() >= 2,
+      "pooled_embs must be at least a 2-D tensor of size [B_local][Sum_T_global(D)], "
       "current shape is: ",
       torch_tensor_shape_str(pooled_embs));
 

fbgemm_gpu/src/permute_pooled_embedding_ops/permute_pooled_embedding_ops_cpu.cpp

@@ -28,8 +28,8 @@ Tensor permute_pooled_embs_cpu_impl(
   }
 
   TORCH_CHECK(
-      pooled_embs.dim() == 2,
-      "pooled_embs must be 2-D tensor of size [B_local][Sum_T_global(D)], "
+      pooled_embs.dim() >= 2,
+      "pooled_embs must be at least a 2-D tensor of size [B_local][Sum_T_global(D)], "
       "current shape is: ",
       torch_tensor_shape_str(pooled_embs));
   TORCH_CHECK(

fbgemm_gpu/src/permute_pooled_embedding_ops/permute_pooled_embedding_ops_split.cu

@@ -70,8 +70,8 @@ Tensor permute_pooled_embs_split_gpu_impl(
   }
 
   TORCH_CHECK(
-      pooled_embs.dim() == 2,
-      "pooled_embs must be 2-D tensor of size [B_local][Sum_T_global(D)], "
+      pooled_embs.dim() >= 2,
+      "pooled_embs must be at least a 2-D tensor of size [B_local][Sum_T_global(D)], "
       "current shape is: ",
       torch_tensor_shape_str(pooled_embs));
 

fbgemm_gpu/src/permute_pooled_embedding_ops/permute_pooled_embedding_ops_split_cpu.cpp

@@ -37,8 +37,8 @@ Tensor permute_pooled_embs_split_cpu_impl(
   }
 
   TORCH_CHECK(
-      pooled_embs.dim() == 2,
-      "pooled_embs must be 2-D tensor of size [B_local][Sum_T_global(D)], "
+      pooled_embs.dim() >= 2,
+      "pooled_embs must be at least a 2-D tensor of size [B_local][Sum_T_global(D)], "
       "current shape is: ",
       torch_tensor_shape_str(pooled_embs));
   TORCH_CHECK(