Solved by MariosOreo in post #4 
              
              
                In your use case, I am afraid that torch.mean() cannot fit our requirement. 
You could have a try on this: 
sum_mask_embeddings = mask_embeddings.sum(dim=1)
for dim in range(mask.size(0)):
    average_mask_embeddings[dim] = sum_mask_embeddings[dim] / mask[dim].sum()