fix: qwen2 lm_head loading huggingface#2443 (huggingface#2445)

Co-authored-by: Yi Xu <[email protected]>
getong · Aug 23, 2024 · fdc2622 · fdc2622
1 parent ccdbe87
commit fdc2622
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/candle-transformers/src/models/qwen2.rs b/candle-transformers/src/models/qwen2.rs
@@ -361,7 +361,7 @@ pub struct ModelForCausalLM {
 impl ModelForCausalLM {
     pub fn new(cfg: &Config, vb: VarBuilder) -> Result<Self> {
         let base_model = Model::new(cfg, vb.clone())?;
-        let lm_head = if vb.contains_tensor("lm_head") {
+        let lm_head = if vb.contains_tensor("lm_head.weight") {
             linear_no_bias(cfg.hidden_size, cfg.vocab_size, vb.pp("lm_head"))?
         } else {
             Linear::from_weights(base_model.embed_tokens.embeddings().clone(), None)