Update handler

handler.py

@@ -3,10 +3,12 @@ import logging
 
 from transformers import AutoModelForCausalLM, AutoTokenizer
 from peft import PeftConfig, PeftModel
+import torch.cuda
 
 
 LOGGER = logging.getLogger(__name__)
 logging.basicConfig(level=logging.INFO)
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
 
 class EndpointHandler():
@@ -29,13 +31,13 @@ class EndpointHandler():
         if prompt is None:
             raise ValueError("Missing prompt.")
         # Preprocess
-        inputs = self.tokenizer(prompt, return_tensors="pt")
+        input_ids = self.tokenizer(prompt, return_tensors="pt").input_ids.to(device)
         # Forward
         LOGGER.info(f"Start generation.")
         if parameters is not None:
-            output = self.model.generate(**inputs, **parameters)
+            output = self.model.generate(input_ids=input_ids, **parameters)
         else:
-            output = self.model.generate(**inputs)
+            output = self.model.generate(input_ids=input_ids)
         # Postprocess
         prediction = self.tokenizer.decode(output[0])
         LOGGER.info(f"Generated text: {prediction}")