Solved by ptrblck in post #2 
              
              
                Yes, that’s the expected output as shown in the docs. 
nn.Embedding layers are adding the embedding_dim dimension as the last additional dimension to the output tensor.