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/opt.py
+++ b/cacheflow/models/opt.py
@@ -234,7 +234,7 @@ class OPTForCausalLM(nn.Module):
        # TODO(zhuohan): create a new weight after implementing pipeline
        #                parallelism
        self.lm_head_weight = self.model.decoder.embed_tokens.weight
-        self.sampler = Sampler()
+        self.sampler = Sampler(config.vocab_size)

    def forward(
        self,