is it better x[-1] @ wte.T

[Sandy4321]

is it better to change return x @ wte.T # [n_seq, n_embd] -> [n_seq, n_vocab] by
x[-1] @ wte.T ?

then we can use next_id = np.argmax(logits)