to 'cuda' instead of model.device
Browse files- handler.py +2 -2
handler.py
CHANGED
@@ -28,7 +28,7 @@ class EndpointHandler():
|
|
28 |
temperature: float = inputs['temperature']
|
29 |
|
30 |
inputs = self.processor(images=image, text=input_text, return_tensors="pt").to(
|
31 |
-
|
32 |
)
|
33 |
output = self.model.generate(
|
34 |
**inputs, max_new_tokens=max_new_tokens, temperature=temperature
|
@@ -47,7 +47,7 @@ class EndpointHandler():
|
|
47 |
|
48 |
inputs = self.processor(
|
49 |
images=image, text=(prompt + continuation), return_tensors="pt"
|
50 |
-
).to(
|
51 |
inputs["labels"] = inputs["input_ids"]
|
52 |
input_ids = inputs["input_ids"][0]
|
53 |
tokens = [self.processor.decode([t]) for t in input_ids]
|
|
|
28 |
temperature: float = inputs['temperature']
|
29 |
|
30 |
inputs = self.processor(images=image, text=input_text, return_tensors="pt").to(
|
31 |
+
'cuda', self.model.dtype
|
32 |
)
|
33 |
output = self.model.generate(
|
34 |
**inputs, max_new_tokens=max_new_tokens, temperature=temperature
|
|
|
47 |
|
48 |
inputs = self.processor(
|
49 |
images=image, text=(prompt + continuation), return_tensors="pt"
|
50 |
+
).to('cuda', self.model.dtype)
|
51 |
inputs["labels"] = inputs["input_ids"]
|
52 |
input_ids = inputs["input_ids"][0]
|
53 |
tokens = [self.processor.decode([t]) for t in input_ids]
|