Update handler.py
Browse files- handler.py +1 -1
handler.py
CHANGED
|
@@ -106,7 +106,7 @@ class EndpointHandler():
|
|
| 106 |
outputs = self.model(**inputs)
|
| 107 |
outputs_logits = outputs.logits[0, -1, self.law_token_ids]
|
| 108 |
|
| 109 |
-
base_input = tokenizer("<cite>", return_tensors="pt").to("cuda")
|
| 110 |
with torch.no_grad():
|
| 111 |
base_output = self.model(**base_input)
|
| 112 |
|
|
|
|
| 106 |
outputs = self.model(**inputs)
|
| 107 |
outputs_logits = outputs.logits[0, -1, self.law_token_ids]
|
| 108 |
|
| 109 |
+
base_input = self.tokenizer("<cite>", return_tensors="pt").to("cuda")
|
| 110 |
with torch.no_grad():
|
| 111 |
base_output = self.model(**base_input)
|
| 112 |
|