Pierce Maloney
commited on
Commit
•
b873ed7
1
Parent(s):
8cda5e7
torch.cuda.empty_cache()
Browse files- handler.py +2 -0
handler.py
CHANGED
|
@@ -1,5 +1,6 @@
|
|
| 1 |
from typing import Dict, List, Any
|
| 2 |
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline, StoppingCriteria, StoppingCriteriaList
|
|
|
|
| 3 |
|
| 4 |
|
| 5 |
class EndpointHandler():
|
|
@@ -18,6 +19,7 @@ class EndpointHandler():
|
|
| 18 |
Return:
|
| 19 |
A :obj:`list` | `dict`: will be serialized and returned
|
| 20 |
"""
|
|
|
|
| 21 |
inputs = data.pop("inputs", data)
|
| 22 |
additional_bad_words_ids = data.pop("additional_bad_words_ids", [])
|
| 23 |
|
|
|
|
| 1 |
from typing import Dict, List, Any
|
| 2 |
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline, StoppingCriteria, StoppingCriteriaList
|
| 3 |
+
import torch
|
| 4 |
|
| 5 |
|
| 6 |
class EndpointHandler():
|
|
|
|
| 19 |
Return:
|
| 20 |
A :obj:`list` | `dict`: will be serialized and returned
|
| 21 |
"""
|
| 22 |
+
torch.cuda.empty_cache()
|
| 23 |
inputs = data.pop("inputs", data)
|
| 24 |
additional_bad_words_ids = data.pop("additional_bad_words_ids", [])
|
| 25 |
|