Fix weights by putting the right value in `lm_head.weight` (#2)

Initial commit