n_layer = 24
n_embd = 1024