Set device=0 for pipeline

handler.py

@@ -8,7 +8,7 @@ class EndpointHandler:
         self.model =  AutoModelForSeq2SeqLM.from_pretrained(path, device_map="auto")
         self.tokenizer = AutoTokenizer.from_pretrained(path)
 
-        self.pipeline = pipeline(task="text-generation", tokenizer=self.tokenizer, device_map="auto", framework="pt", model=self.model, max_length=512)
+        self.pipeline = pipeline(task="text-generation", tokenizer=self.tokenizer, device=0, device_map="auto", framework="pt", model=self.model, max_length=512)
 
     def __call__(self, data: Dict[str, Any]) -> Dict[str, str]:
         """
@@ -25,9 +25,9 @@ class EndpointHandler:
 
         # pass inputs with all kwargs in data
         if parameters is not None:
-            prediction = self.pipeline(inputs, **parameters)
+            prediction = self.pipeline(inputs, device=0, **parameters)
         else:
-            prediction = self.pipeline(inputs)
+            prediction = self.pipeline(inputs, device=0)
 
         # postprocess the prediction
         prediction = self.tokenizer.decode(outputs[0], skip_special_tokens=True)