Train nn.embedding

[ptrblck]

Double post from here.
Let’s continue the discussion in the other thread, please.