Thank you for the suggestion. I have now tried to do the following:

Initialize a weight tensor and
adjust the forward function accordingly. But now I run into memory problems (see end of post):

Initializing weights:
group_ids = d.keys()
weights_shape = (len(group_ids), P.shape[0])
weights = torch.zeros(size=weights_shape, dtype=torch.double, requires_grad=False, device='cuda')

for group_id in group_ids:
    number_of_relevant_embs = len(d[group_id])
    id_of_first_relevant_emb = d[group_id][0]
    id_of_last_relevant_emb = id_of_first_relevant_emb+number_of_relevant_embs
    weights[group_id][id_of_first_relevant_emb:id_of_last_relevant_emb] = 1/number_of_relevant_embs

Adjusted part of forward function (given batch of group_ids):
weight_batch = weights[batch]
batch_of_averaged_embeddings = torch.matmul(weight_batch, embeddings)

Unfortunately, since weights_shape is around (30000,40000), this approach already fails when trying to create the weights torch.zeros tensor, with a CUDA out of memory RuntimeError (wants to allocate around 10GB for this, which is almost the entire memory of my GPU).
Is what I did what you meant in the first place?