inference time mini-optimization low-hanging fruit ty @jxtps for raising: when we are running inference we can apply lm_head on only the very last token

train.py

@@ -102,7 +102,7 @@ def get_batch(split):
 iter_num = 0
 best_val_loss = 1e9
 
-# model init
+# model init. TODO: fix bug we should also propagate the correct vocab_size to the model_args
 model_args = dict(n_layer = n_layer, n_head = n_head, n_embd = n_embd, block_size = block_size, dropout = dropout)
 if init_from == 'scratch':
     # init a new model from scratch