Update handler.py
Browse files- handler.py +2 -0
handler.py
CHANGED
@@ -36,6 +36,8 @@ class EndpointHandler:
|
|
36 |
"", # output - leave this blank for generation!
|
37 |
)
|
38 |
], return_tensors="pt")
|
|
|
|
|
39 |
|
40 |
outputs = self.model.generate(**inputs,
|
41 |
max_new_tokens=1048,
|
|
|
36 |
"", # output - leave this blank for generation!
|
37 |
)
|
38 |
], return_tensors="pt")
|
39 |
+
|
40 |
+
inputs = inputs.to('cuda')
|
41 |
|
42 |
outputs = self.model.generate(**inputs,
|
43 |
max_new_tokens=1048,
|