Update aquila.py

FlagAI-Open · Aug 2, 2023 · d3722a6 · d3722a6
1 parent 7aeb339
commit d3722a6
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/flagai/model/predictor/aquila.py b/flagai/model/predictor/aquila.py
@@ -29,8 +29,8 @@ def aquila_generate(
 
         total_len = min(2048, max_gen_len + max_prompt_size)
 
-        # tokens = torch.full((bsz, total_len), 0).cuda().long()
-        tokens = torch.full((bsz, total_len), 0).to("cuda:5").long()
+        tokens = torch.full((bsz, total_len), 0).cuda().long()
+        #tokens = torch.full((bsz, total_len), 0).to("cuda:5").long()
         for k, t in enumerate(prompt_tokens):
             tokens[k, : len(t)] = t.clone().detach().long()
         input_text_mask = tokens != 0