[shape inference] fix pruning
Summary: Use the dim type of the first input for output.
Test Plan:
unit test
flow test: f254777437
https://fburl.com/n933wc3a
shapes {
shape {
dims: 19102004
dims: 68
data_type: UINT8
name: "sparse_nn_2/sparse_arch_2/grouped_embedding_10/grouped_generic_embedding_10/GSF_IDLIST_IG_BUSINESS_AUTHOR_PPR_ORGANIC_ENGAGEMENT_UNIFORM_RIDS/w_EmbeddingFusedUint4Quantization"
}
dim_type: CONSTANT
dim_type: CONSTANT
name: "sparse_nn_2/sparse_arch_2/grouped_embedding_10/grouped_generic_embedding_10/GSF_IDLIST_IG_BUSINESS_AUTHOR_PPR_ORGANIC_ENGAGEMENT_UNIFORM_RIDS/w_EmbeddingFusedUint4Quantization"
shape_is_final: true
}
Reviewed By: yinghai, khabinov
Differential Revision: D26763978
fbshipit-source-id: b9c0d6ca4a2b0e4d50d34e08f724e99ad705196b