So far, I have the following:
encoder = nn.Embedding(vocab_size, emb_size)
l = []
for i in range(seq_length * batch_size):
    l.append(encoder.weight)
t = torch.stack(l).view(seq_length, batch_size, vocab_size, emb_size)

I’m not sure I can reuse the embedding matrix like this and also even if this is correct, I really feel like there has to be a better way of doing this