fix row-wise alltoall error when some embeddings use mean pooling and others use sum pooling (#2809)

Summary:
KJT.lengths is modified by mean pooling callback.  When some embeddings use mean pooling and others use sum pooling, KJT.lengths will be incorrect.

Pull Request resolved: #2809

Reviewed By: PaulZhang12

Differential Revision: D71759505

Pulled By: iamzainhuda

fbshipit-source-id: d899a2a072f84e9b6692aac0653914f5d3490f27

Author: tiankongdeguiji

torchrec/distributed/embeddingbag.py

@@ -1716,6 +1716,7 @@ def _create_mean_pooling_divisor(
             lengths = torch.index_select(input=lengths, dim=0, index=indices)
 
         # only convert the sum pooling features to be 1 lengths
+        lengths = lengths.clone()
         for feature in pooling_type_to_rs_features[PoolingType.SUM.value]:
             feature_index = kjt_key_indices[feature]
             feature_index = feature_index * batch_size