pub fn decode_head(spec: &ModelSpec, dir: &str) -> Authored
A decode step’s head alone: the final norm, the LM head and the arg-max.