I’ve also discovered that multiplying embedding matrix with onehot vectors is much faster than F.embedding, but creating onehot vectors with F.onehot is exteremly slow. I created a “whiteboard” vector filled with zeros, then scattered the indices onto it to create onehot vectors, this is much faster than creating a new tensor at each training step.
batch_size=128
seq_len =100
vocab_size=10000
w = torch.rand(vocab_size, 256, requires_grad=True)
whiteboard = torch.zeros([batch_size, seq_len, vocab_size])

for i in range(100):
  x = torch.randint(vocab_size, size=[batch_size, seq_len])
  onehot = whiteboard.scatter(2, x[:,:,None], 1)
  # onehot = F.onehot(x, vocab_size) # very slow
  
  y = onehot @ w
  # y = F.embedding(x, w)

  # ...