apoorvkh commited on
Commit
007e4d7
1 Parent(s): ca0693b

to 'cuda' instead of model.device

Browse files
Files changed (1) hide show
  1. 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
- self.model.device, self.model.dtype
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(self.model.device, 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]
 
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]