Use high priority stream for forward pass

test_inference.py

@@ -590,6 +590,7 @@ if args.speed:
 
                 logits = model.forward(ids[:, -1:], cache)
                 sample = torch.argmax(logits[0, -1]).cpu().unsqueeze(0).unsqueeze(0)
+                sample.clamp_(0, tokenizer.get_vocab_size() - 1)
                 ids = torch.cat((ids, sample), dim=-1)
 
             time_end = time.time()