I have a case where I get a smoothened one-hot(probability) distribution and I would like to do an average embedding lookup. This is what I am doing now,
    embedding = nn.Embedding(vocab_size, embdim)
    if flag:
        embedded = embedding(input_seqs)
    else:
        embedded = torch.matmul(input_seqs, embedding.weight.data)

when flag = True I will be passing normal input sequence as indices of size (batch x max_seq_len), while in other case it will be a smoothened onehot vectors of size (batch x max_seq_len x vocab_size).
Is this the correct way to do it?
Thanks