Update gpt2 to use wte if no lm_head

crates/models/gpt2/src/lib.rs

@@ -31,7 +31,7 @@ pub struct Gpt2 {
     // weighted positional encodings
     wpe: Tensor,
     // language model head
-    lm_head: Tensor,
+    lm_head: Option<Tensor>,
 
     // weights for the model
     layers: Vec<Layer>,
@@ -59,7 +59,7 @@ impl KnownModel for Gpt2 {
         let ln_f_b = tl.load("model/ln_f/b")?;
         let wte = tl.load("model/wte")?;
         let wpe = tl.load("model/wpe")?;
-        let lm_head = tl.load("model/lm_head")?;
+        let lm_head = tl.load("model/lm_head").ok();
 
         let mut layers = Vec::new();
         for i in 0..hyperparameters.n_layer {
@@ -306,7 +306,8 @@ impl KnownModel for Gpt2 {
 
             let embeddings_tensor: ggml::Tensor = input_layer.share();
 
-            input_layer = ctx0.op_mul_mat(&self.lm_head, &input_layer);
+            let head = self.lm_head.as_ref().unwrap_or(&self.wte);
+            input_layer = ctx0.op_mul_mat(head, &input_layer);
 
             (
                 gf,