aaabiao commited on
Commit
57908e7
1 Parent(s): 4b0cfea

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +1 -1
app.py CHANGED
@@ -16,7 +16,7 @@ from transformers import (
16
  class StoppingCriteriaSub(StoppingCriteria):
17
  def __init__(self, stops = [], encounters=1):
18
  super().__init__()
19
- # self.stops = [stop.to("cuda") for stop in stops]
20
 
21
  def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor):
22
  last_token = input_ids[0][-1]
 
16
  class StoppingCriteriaSub(StoppingCriteria):
17
  def __init__(self, stops = [], encounters=1):
18
  super().__init__()
19
+ self.stops = [stop.to("cuda") for stop in stops]
20
 
21
  def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor):
22
  last_token = input_ids[0][-1]