I’m not sure to understand the question completely, so please correct me, if I’m misunderstanding something.
If you would like to finetune an embedding layer, you can use the from_pretrained method and set freese=False, as you suggested.
The shapes should work out of the box:
# Create embedding
N, D = 10, 100
emb = nn.Embedding(num_embeddings=N, embedding_dim=D)

# Create input
batch_size = 5
L = 7
x = torch.randint(0, N, (batch_size, L))

# Forward pass
output = emb(x)
print(output.shape)  # batch_size, L, D