Spaces:
Running
on
Zero
Running
on
Zero
Update app.py
Browse files
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 |
-
|
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]
|