Weighted mean using nn.EmbeddingBag

I want to calculate weighted mean of ‘bags’ of embeddings rather than the simple mean. Is there a way I can do this efficiently. As far as I could see currently nn.EmbeddingBag does not support this functionality, is someone working on it?