reshape inputs for GPU

handler.py

@@ -10,7 +10,7 @@ class EndpointHandler():
 
     def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
         input = data.get("inputs", data)
-        inputs = self.tokenizer(input, return_tensors="pt", return_attention_mask=False)
+        inputs = self.tokenizer(input, return_tensors="pt", return_attention_mask=False).to('cuda')
         outputs = self.model.generate(**inputs, max_length=len(inputs[0]) + 8, pad_token_id=self.tokenizer.eos_token_id)
         text = self.tokenizer.batch_decode(outputs)[0]
         return text