run on cpu

3 years ago · e04d724890
parent b998e985a6
commit e04d724890
1 changed files with 6 additions and 6 deletions
--- a/llama/generation.py
+++ b/llama/generation.py
@ -32,7 +32,7 @@ class LLaMA:

        total_len = min(params.max_seq_len, max_gen_len + max_prompt_size)

-        tokens = torch.full((bsz, total_len), self.tokenizer.pad_id).cuda().long()
+        tokens = torch.full((bsz, total_len), self.tokenizer.pad_id).cpu().long()
        for k, t in enumerate(prompt_tokens):
            tokens[k, : len(t)] = torch.tensor(t).long()
        input_text_mask = tokens != self.tokenizer.pad_id