inputs to device 0 instead of cuda

handler.py

@@ -28,7 +28,7 @@ class EndpointHandler():
             temperature: float = inputs['temperature']
 
             inputs = self.processor(images=image, text=input_text, return_tensors="pt").to(
-                'cuda', self.model.dtype
+                0, self.model.dtype
             )
             output = self.model.generate(
                 **inputs, max_new_tokens=max_new_tokens, temperature=temperature
@@ -47,7 +47,7 @@ class EndpointHandler():
 
             inputs = self.processor(
                 images=image, text=(prompt + continuation), return_tensors="pt"
-            ).to('cuda', self.model.dtype)
+            ).to(0, self.model.dtype)
             inputs["labels"] = inputs["input_ids"]
             input_ids = inputs["input_ids"][0]
             tokens = [self.processor.decode([t]) for t in input_ids]