Add support for GPT-2 (#60)

2023-05-04 02:59:56 -07:00
parent 130d5fd8c7
commit e548c1488a
7 changed files with 350 additions and 8 deletions
--- a/cacheflow/models/llama.py
+++ b/cacheflow/models/llama.py
@@ -192,7 +192,7 @@ class LlamaForCausalLM(nn.Module):
                                            bias=False,
                                            gather_output=False,
                                            perform_initialization=False)
-        self.sampler = Sampler()
+        self.sampler = Sampler(config.vocab_size)

    def forward(
        self,