Hello @MariosOreo, first, thanks for the answer.
It works, but in the mask_embeddings.mean(dim=1) you gonna consider the vectors in the B dimension that are 0 in the mask for the mean… Although they don’t contribute for the sum, it is going to contribute as a value for the division part of the mean, right?
In the example I gave:

embeddings.size() -> torch.Size([2, 3, 1024])
mask.size() -> torch.Size([2, 3])
mask -> tensor([[1, 1, 0], [1, 1, 1]])

The example 0 of the batch is going to have a size of torch.Size([2, 1024]). Whereas the example 1 of the batch torch.Size([3, 1024]), and I would like to take the mean for the example 0 over the 2 word embeddings, and over the 3 word embeddings for the example 1.