set the device id for the taskflow (PaddlePaddle#1011)

Author: wawltor

paddlenlp/taskflow/task.py

@@ -91,7 +91,7 @@ def _prepare_static_mode(self):
         if place == 'cpu':
             self._config.disable_gpu()
         else:
-            self._config.enable_use_gpu(100, 0)
+            self._config.enable_use_gpu(100, self.kwargs['device_id'])
         self._config.switch_use_feed_fetch_ops(False)
         self._config.disable_glog_info()
         self.predictor = paddle.inference.create_predictor(self._config)

paddlenlp/taskflow/taskflow.py

@@ -135,6 +135,7 @@ def __init__(self, task, model=None, device_id=0, **kwargs):
         self.model = model
         # Update the task config to kwargs
         config_kwargs = TASKS[self.task]['models'][self.model]
+        kwargs['device_id'] = device_id
         kwargs.update(config_kwargs)
         self.kwargs = kwargs
         task_class = TASKS[self.task]['models'][self.model]['task_class']