Repository: huggingface/optimum

Issue:
Not possible to configure GPU in pipelines nor leveraging batch_size parallelisation
When setting the `device` variable in the `pipeline` function/class to `>= 0`, an error appears `AttributeError: 'ORTModelForCausalLM' object has no attribute 'to' - when running in GPU`. This was initially reported in #161 so opening this issue to encompass supporting the `device` parameter in the ORT classes. This is important as otherwise it won't be possible to allow configuration of CPU/GPU similar to normal transformer libraries.

Is there currently a workaround to ensure that the class is run on GPU? By default it seems this woudl eb set in CPU even when GPU is available:
```python
>>> m = ORTModelForCausalLM.from_pretrained("gpt2", from_transformers=True)
>>> t = AutoTokenizer.from_pretrained("gpt2")
>>> pp = pipeline("text-generation", model=m, tokenizer=t)
>>> pp.device

device(type='cpu')
```

This is still the case even with the `optimum[onnxruntime-gpu]` package. I have validated by testing against a normal transformer with `batch_size=X` (ie `pp = pipeline("text-generation", model=m, tokenizer=t, batch_size=128)`) and it seems there is no optimization with parallel processing with optimum, whereas normal transformer is orders of magnitude faster (which is most likely as it's not utilizing the parallelism)

I can confirm that the model is loaded with GPU correctly:

```python
>>> m.device

device(type='cuda', index=0)
```

And GPU is configured correctly:

```python
>>> from optimum.onnxruntime.utils import _is_gpu_available
>>> _is_gpu_available()

True
```

Is there a way to enable GPU for processing with batching in optimum?
Passing GPU Device to ORT Pipeline fails
I was trying to run a GPU accelerated optimum pipeline but passing device to the pipeline object fails in the following way

```python
from optimum.pipelines import pipeline as ort_pipeline

ort_pipeline(
  "text-classification", model=quantized_model, tokenizer=tokenizer, return_all_scores=True, device=0
)
```

returns error
`AttributeError: 'ORTModelForSequenceClassification' object has no attribute 'to'`

This works for the standard transformers pipeline object. I am exporting my quantized model which does not seem to support a `to_gpu` flag like the fit command does. Wondering how people export a quantized model then run it on GPU.



Discussion:

Resolving patch:
diff --git a/optimum/onnxruntime/modeling_ort.py b/optimum/onnxruntime/modeling_ort.py
--- a/optimum/onnxruntime/modeling_ort.py
+++ b/optimum/onnxruntime/modeling_ort.py
@@ -30,7 +30,7 @@
 from huggingface_hub import HfApi, hf_hub_download
 
 from ..modeling_base import OptimizedModel
-from .utils import ONNX_WEIGHTS_NAME, _is_gpu_available
+from .utils import ONNX_WEIGHTS_NAME, get_device_for_provider, get_provider_for_device
 
 
 logger = logging.getLogger(__name__)
@@ -85,28 +85,50 @@ def __init__(self, model=None, config=None, **kwargs):
         self.config = config
         self.model_save_dir = kwargs.get("model_save_dir", None)
         self.latest_model_name = kwargs.get("latest_model_name", "model.onnx")
+        self._device = get_device_for_provider(self.model.get_providers()[0])
 
         # registers the ORTModelForXXX classes into the transformers AutoModel classes
         # to avoid warnings when create a pipeline https://github.com/huggingface/transformers/blob/cad61b68396a1a387287a8e2e2fef78a25b79383/src/transformers/pipelines/base.py#L863
         AutoConfig.register(self.base_model_prefix, AutoConfig)
         self.auto_model_class.register(AutoConfig, self.__class__)
 
+    @property
+    def device(self) -> torch.device:
+        """
+        `torch.device`: The device on which the module is (assuming that all the module parameters are on the same
+        device).
+        """
+        return self._device
+
+    @device.setter
+    def device(self, value):
+        self._device = value
+
+    def to(self, device):
+        """
+        Changes the ONNX Runtime provider according to the device.
+        """
+        self.device = device
+        provider = get_provider_for_device(self.device)
+        self.model.set_providers([provider])
+        return self
+
     def forward(self, *args, **kwargs):
         raise NotImplementedError
 
     @staticmethod
     def load_model(path: Union[str, Path], provider=None):
         """
-        loads ONNX Inference session with Provider. Default Provider is if CUDAExecutionProvider GPU available else `CPUExecutionProvider`
+        Loads an ONNX Inference session with a given provider. Default provider is `CPUExecutionProvider` to match the default behaviour in PyTorch/TensorFlow/JAX.
+
         Arguments:
             path (`str` or `Path`):
-                Directory from which to load
+                Directory from which to load the model.
             provider(`str`, *optional*):
-                Onnxruntime provider to use for loading the model, defaults to `CUDAExecutionProvider` if GPU is
-                available else `CPUExecutionProvider`
+                ONNX Runtime provider to use for loading the model. Defaults to `CPUExecutionProvider`.
         """
         if provider is None:
-            provider = "CUDAExecutionProvider" if _is_gpu_available() else "CPUExecutionProvider"
+            provider = "CPUExecutionProvider"
 
         return ort.InferenceSession(path, providers=[provider])
 
@@ -330,10 +352,9 @@ def forward(
             onnx_inputs["token_type_ids"] = token_type_ids.cpu().detach().numpy()
         # run inference
         outputs = self.model.run(None, onnx_inputs)
+        last_hidden_state = torch.from_numpy(outputs[self.model_outputs["last_hidden_state"]]).to(self.device)
         # converts output to namedtuple for pipelines post-processing
-        return BaseModelOutput(
-            last_hidden_state=torch.from_numpy(outputs[self.model_outputs["last_hidden_state"]]),
-        )
+        return BaseModelOutput(last_hidden_state=last_hidden_state)
 
 
 QUESTION_ANSWERING_SAMPLE = r"""
@@ -416,10 +437,12 @@ def forward(
             onnx_inputs["token_type_ids"] = token_type_ids.cpu().detach().numpy()
         # run inference
         outputs = self.model.run(None, onnx_inputs)
+        start_logits = torch.from_numpy(outputs[self.model_outputs["start_logits"]]).to(self.device)
+        end_logits = torch.from_numpy(outputs[self.model_outputs["end_logits"]]).to(self.device)
         # converts output to namedtuple for pipelines post-processing
         return QuestionAnsweringModelOutput(
-            start_logits=torch.from_numpy(outputs[self.model_outputs["start_logits"]]),
-            end_logits=torch.from_numpy(outputs[self.model_outputs["end_logits"]]),
+            start_logits=start_logits,
+            end_logits=end_logits,
         )
 
 
@@ -519,9 +542,10 @@ def forward(
             onnx_inputs["token_type_ids"] = token_type_ids.cpu().detach().numpy()
         # run inference
         outputs = self.model.run(None, onnx_inputs)
+        logits = torch.from_numpy(outputs[self.model_outputs["logits"]]).to(self.device)
         # converts output to namedtuple for pipelines post-processing
         return SequenceClassifierOutput(
-            logits=torch.from_numpy(outputs[self.model_outputs["logits"]]),
+            logits=logits,
         )
 
 
@@ -604,9 +628,10 @@ def forward(
             onnx_inputs["token_type_ids"] = token_type_ids.cpu().detach().numpy()
         # run inference
         outputs = self.model.run(None, onnx_inputs)
+        logits = torch.from_numpy(outputs[self.model_outputs["logits"]]).to(self.device)
         # converts output to namedtuple for pipelines post-processing
         return TokenClassifierOutput(
-            logits=torch.from_numpy(outputs[self.model_outputs["logits"]]),
+            logits=logits,
         )
 
 
@@ -665,14 +690,6 @@ def __init__(self, *args, **kwargs):
         self.main_input_name = "input_ids"
         self.model_outputs = {output_key.name: idx for idx, output_key in enumerate(self.model.get_outputs())}
 
-    @property
-    def device(self) -> torch.device:
-        """
-        `torch.device`: The device on which the module is (assuming that all the module parameters are on the same
-        device).
-        """
-        return torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-
     def prepare_inputs_for_generation(self, input_ids: torch.LongTensor, **kwargs) -> Dict[str, Any]:
         """
         Implement in subclasses of [`PreTrainedModel`] for custom behavior to prepare inputs in the generate method.
@@ -703,9 +720,10 @@ def forward(
         }
         # run inference
         outputs = self.model.run(None, onnx_inputs)
+        logits = torch.from_numpy(outputs[self.model_outputs["logits"]]).to(self.device)
         # converts output to namedtuple for pipelines post-processing
         return CausalLMOutputWithCrossAttentions(
-            logits=torch.from_numpy(outputs[self.model_outputs["logits"]]),
+            logits=logits,
         )
 
     # Adapted from https://github.com/huggingface/transformers/blob/99289c08a1b16a805dd4ee46de029e9fd23cba3d/src/transformers/generation_utils.py#L490
diff --git a/optimum/onnxruntime/utils.py b/optimum/onnxruntime/utils.py
--- a/optimum/onnxruntime/utils.py
+++ b/optimum/onnxruntime/utils.py
@@ -141,3 +141,17 @@ def wrap_onnx_config_for_loss(onnx_config: OnnxConfig) -> OnnxConfig:
         return OnnxConfigWithPastAndLoss(onnx_config)
     else:
         return OnnxConfigWithLoss(onnx_config)
+
+
+def get_device_for_provider(provider: str) -> torch.device:
+    """
+    Gets the PyTorch device (CPU/CUDA) associated with an ONNX Runtime provider.
+    """
+    return torch.device("cuda") if provider == "CUDAExecutionProvider" else torch.device("cpu")
+
+
+def get_provider_for_device(device: torch.device) -> str:
+    """
+    Gets the ONNX Runtime provider associated with the PyTorch device (CPU/CUDA).
+    """
+    return "CUDAExecutionProvider" if device.type.lower() == "cuda" else "CPUExecutionProvider"


Test patch:
diff --git a/.github/workflows/test_onnxruntime_gpu.yml b/.github/workflows/test_onnxruntime_gpu.yml
new file mode 100644
--- /dev/null
+++ b/.github/workflows/test_onnxruntime_gpu.yml
@@ -0,0 +1,85 @@
+name: ONNX Runtime / Test GPU
+
+on:
+  workflow_dispatch:
+  schedule:
+    - cron: 0 7 * * * # every day at 7am
+
+jobs:
+  start-runner:
+    name: Start self-hosted EC2 runner
+    runs-on: ubuntu-latest
+    env:
+      AWS_REGION: us-east-1
+      EC2_AMI_ID: ami-0dc1c26161f869ed1
+      EC2_INSTANCE_TYPE: g4dn.xlarge
+      EC2_SUBNET_ID: subnet-859322b4,subnet-b7533b96,subnet-47cfad21,subnet-a396b2ad,subnet-06576a4b,subnet-df0f6180
+      EC2_SECURITY_GROUP: sg-0bb210cd3ec725a13
+      EC2_IAM_ROLE: optimum-ec2-github-actions-role
+    outputs:
+      label: ${{ steps.start-ec2-runner.outputs.label }}
+      ec2-instance-id: ${{ steps.start-ec2-runner.outputs.ec2-instance-id }}
+    steps:
+      - name: Configure AWS credentials
+        uses: aws-actions/configure-aws-credentials@v1
+        with:
+          aws-access-key-id: ${{ secrets.AWS_ACCESS_KEY_ID }}
+          aws-secret-access-key: ${{ secrets.AWS_SECRET_ACCESS_KEY }}
+          aws-region: ${{ env.AWS_REGION }}
+      - name: Start EC2 runner
+        id: start-ec2-runner
+        uses: philschmid/philschmid-ec2-github-runner@main
+        with:
+          mode: start
+          github-token: ${{ secrets.GH_PERSONAL_ACCESS_TOKEN }}
+          ec2-image-id: ${{ env.EC2_AMI_ID }}
+          ec2-instance-type: ${{ env.EC2_INSTANCE_TYPE }}
+          subnet-id: ${{ env.EC2_SUBNET_ID }}
+          security-group-id: ${{ env.EC2_SECURITY_GROUP }}
+          iam-role-name: ${{ env.EC2_IAM_ROLE }}
+          aws-resource-tags: > # optional, requires additional permissions
+            [
+              {"Key": "Name", "Value": "ec2-optimum-github-runner"},
+              {"Key": "GitHubRepository", "Value": "${{ github.repository }}"}
+            ]
+  do-the-job:
+    name: Setup
+    needs: start-runner # required to start the main job when the runner is ready
+    runs-on: ${{ needs.start-runner.outputs.label }} # run the job on the newly created runner
+    env:
+      AWS_REGION: us-east-1
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v2
+      - name: Install dependencies
+        run: |
+          sudo apt -y update && sudo pip install --upgrade pip
+          pip install .[onnxruntime-gpu,tests]
+    - name: Test with unittest
+      working-directory: tests
+      run: |
+        python -m unittest discover -s onnxruntime -p 'test_*.py'
+
+  stop-runner:
+    name: Stop self-hosted EC2 runner
+    needs:
+      - start-runner # required to get output from the start-runner job
+      - do-the-job # required to wait when the main job is done
+    runs-on: ubuntu-latest
+    env:
+      AWS_REGION: us-east-1
+    if: ${{ always() }} # required to stop the runner even if the error happened in the previous jobs
+    steps:
+      - name: Configure AWS credentials
+        uses: aws-actions/configure-aws-credentials@v1
+        with:
+          aws-access-key-id: ${{ secrets.AWS_ACCESS_KEY_ID }}
+          aws-secret-access-key: ${{ secrets.AWS_SECRET_ACCESS_KEY }}
+          aws-region: ${{ env.AWS_REGION }}
+      - name: Stop EC2 runner
+        uses: philschmid/philschmid-ec2-github-runner@main
+        with:
+          mode: stop
+          github-token: ${{ secrets.GH_PERSONAL_ACCESS_TOKEN }}
+          label: ${{ needs.start-runner.outputs.label }}
+          ec2-instance-id: ${{ needs.start-runner.outputs.ec2-instance-id }}
diff --git a/tests/onnxruntime/test_modeling_ort.py b/tests/onnxruntime/test_modeling_ort.py
--- a/tests/onnxruntime/test_modeling_ort.py
+++ b/tests/onnxruntime/test_modeling_ort.py
@@ -15,6 +15,7 @@
     PretrainedConfig,
     pipeline,
 )
+from transformers.testing_utils import require_torch_gpu
 
 import onnxruntime
 from optimum.onnxruntime import (
@@ -31,7 +32,7 @@
 from parameterized import parameterized
 
 
-class ORTModelIntergrationTest(unittest.TestCase):
+class ORTModelIntegrationTest(unittest.TestCase):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.TEST_MODEL_ID = "sshleifer/tiny-distilbert-base-cased-distilled-squad"
@@ -54,6 +55,19 @@ def test_load_model_from_hub_without_onnx_model(self):
             ORTModel.from_pretrained(self.FAIL_ONNX_MODEL_ID)
         self.assertEqual("Not Found", context.exception.response.reason)
 
+    def test_model_on_cpu(self):
+        model = ORTModel.from_pretrained(self.ONNX_MODEL_ID)
+        cpu = torch.device("cpu")
+        model.to(cpu)
+        self.assertEqual(model.device, cpu)
+
+    @require_torch_gpu
+    def test_model_on_gpu(self):
+        model = ORTModel.from_pretrained(self.ONNX_MODEL_ID)
+        gpu = torch.device("cuda")
+        model.to(gpu)
+        self.assertEqual(model.device, gpu)
+
     @require_hf_token
     def test_load_model_from_hub_private(self):
         model = ORTModel.from_pretrained(self.ONNX_MODEL_ID, use_auth_token=os.environ.get("HF_AUTH_TOKEN", None))
@@ -94,7 +108,7 @@ def test_save_model_from_hub(self):
             )
 
 
-class ORTModelForQuestionAnsweringIntergrationTest(unittest.TestCase):
+class ORTModelForQuestionAnsweringIntegrationTest(unittest.TestCase):
     SUPPORTED_ARCHITECTURES_WITH_MODEL_ID = {
         "distilbert": "hf-internal-testing/tiny-random-distilbert",
         "bert": "hf-internal-testing/tiny-random-bert",
@@ -172,8 +186,32 @@ def test_pipeline(self, *args, **kwargs):
         self.assertGreaterEqual(outputs["score"], 0.0)
         self.assertTrue(isinstance(outputs["answer"], str))
 
+    @parameterized.expand(SUPPORTED_ARCHITECTURES_WITH_MODEL_ID.items())
+    @require_torch_gpu
+    def test_pipeline_on_gpu(self, *args, **kwargs):
+        model_arch, model_id = args
+        onnx_model = ORTModelForQuestionAnswering.from_pretrained(model_id, from_transformers=True)
+        tokenizer = AutoTokenizer.from_pretrained(model_id)
+        pp = pipeline("question-answering", model=onnx_model, tokenizer=tokenizer, device=0)
+        question = "Whats my name?"
+        context = "My Name is Philipp and I live in Nuremberg."
+        outputs = pp(question, context)
+        # check model device
+        self.assertEqual(pp.model.device.type.lower(), "cuda")
+        # compare model output class
+        self.assertGreaterEqual(outputs["score"], 0.0)
+        self.assertTrue(isinstance(outputs["answer"], str))
+
+    @parameterized.expand(SUPPORTED_ARCHITECTURES_WITH_MODEL_ID.items())
+    def test_default_pipeline_and_model_device(self, *args, **kwargs):
+        model_arch, model_id = args
+        onnx_model = ORTModelForQuestionAnswering.from_pretrained(model_id, from_transformers=True)
+        tokenizer = AutoTokenizer.from_pretrained(model_id)
+        pp = pipeline("question-answering", model=onnx_model, tokenizer=tokenizer)
+        self.assertEqual(pp.device, pp.model.device)
+
 
-class ORTModelForSequenceClassificationIntergrationTest(unittest.TestCase):
+class ORTModelForSequenceClassificationIntegrationTest(unittest.TestCase):
     SUPPORTED_ARCHITECTURES_WITH_MODEL_ID = {
         "distilbert": "hf-internal-testing/tiny-random-distilbert",
         "bert": "hf-internal-testing/tiny-random-bert",
@@ -246,6 +284,29 @@ def test_pipeline(self, *args, **kwargs):
         self.assertGreaterEqual(outputs[0]["score"], 0.0)
         self.assertTrue(isinstance(outputs[0]["label"], str))
 
+    @parameterized.expand(SUPPORTED_ARCHITECTURES_WITH_MODEL_ID.items())
+    @require_torch_gpu
+    def test_pipeline_on_gpu(self, *args, **kwargs):
+        model_arch, model_id = args
+        onnx_model = ORTModelForSequenceClassification.from_pretrained(model_id, from_transformers=True)
+        tokenizer = AutoTokenizer.from_pretrained(model_id)
+        pp = pipeline("text-classification", model=onnx_model, tokenizer=tokenizer, device=0)
+        text = "My Name is Philipp and i live in Germany."
+        outputs = pp(text)
+        # check model device
+        self.assertEqual(pp.model.device.type.lower(), "cuda")
+        # compare model output class
+        self.assertGreaterEqual(outputs[0]["score"], 0.0)
+        self.assertTrue(isinstance(outputs[0]["label"], str))
+
+    @parameterized.expand(SUPPORTED_ARCHITECTURES_WITH_MODEL_ID.items())
+    def test_default_pipeline_and_model_device(self, *args, **kwargs):
+        model_arch, model_id = args
+        onnx_model = ORTModelForSequenceClassification.from_pretrained(model_id, from_transformers=True)
+        tokenizer = AutoTokenizer.from_pretrained(model_id)
+        pp = pipeline("text-classification", model=onnx_model, tokenizer=tokenizer)
+        self.assertEqual(pp.device, onnx_model.device)
+
     def test_pipeline_zero_shot_classification(self):
         onnx_model = ORTModelForSequenceClassification.from_pretrained(
             "typeform/distilbert-base-uncased-mnli", from_transformers=True
@@ -262,7 +323,7 @@ def test_pipeline_zero_shot_classification(self):
         self.assertTrue(any(isinstance(label, str) for label in outputs["labels"]))
 
 
-class ORTModelForTokenClassificationIntergrationTest(unittest.TestCase):
+class ORTModelForTokenClassificationIntegrationTest(unittest.TestCase):
     SUPPORTED_ARCHITECTURES_WITH_MODEL_ID = {
         "distilbert": "hf-internal-testing/tiny-random-distilbert",
         "bert": "hf-internal-testing/tiny-random-bert",
@@ -332,8 +393,30 @@ def test_pipeline(self, *args, **kwargs):
         # compare model output class
         self.assertTrue(any(item["score"] > 0.0 for item in outputs))
 
+    @parameterized.expand(SUPPORTED_ARCHITECTURES_WITH_MODEL_ID.items())
+    @require_torch_gpu
+    def test_pipeline_on_gpu(self, *args, **kwargs):
+        model_arch, model_id = args
+        onnx_model = ORTModelForTokenClassification.from_pretrained(model_id, from_transformers=True)
+        tokenizer = AutoTokenizer.from_pretrained(model_id)
+        pp = pipeline("token-classification", model=onnx_model, tokenizer=tokenizer, device=0)
+        text = "My Name is Philipp and i live in Germany."
+        outputs = pp(text)
+        # check model device
+        self.assertEqual(pp.model.device.type.lower(), "cuda")
+        # compare model output class
+        self.assertTrue(any(item["score"] > 0.0 for item in outputs))
+
+    @parameterized.expand(SUPPORTED_ARCHITECTURES_WITH_MODEL_ID.items())
+    def test_default_pipeline_and_model_device(self, *args, **kwargs):
+        model_arch, model_id = args
+        onnx_model = ORTModelForTokenClassification.from_pretrained(model_id, from_transformers=True)
+        tokenizer = AutoTokenizer.from_pretrained(model_id)
+        pp = pipeline("token-classification", model=onnx_model, tokenizer=tokenizer)
+        self.assertEqual(pp.device, onnx_model.device)
 
-class ORTModelForFeatureExtractionIntergrationTest(unittest.TestCase):
+
+class ORTModelForFeatureExtractionIntegrationTest(unittest.TestCase):
     SUPPORTED_ARCHITECTURES_WITH_MODEL_ID = {
         "distilbert": "hf-internal-testing/tiny-random-distilbert",
         "bert": "hf-internal-testing/tiny-random-bert",
@@ -403,8 +486,30 @@ def test_pipeline(self, *args, **kwargs):
         # compare model output class
         self.assertTrue(any(any(isinstance(item, float) for item in row) for row in outputs[0]))
 
+    @parameterized.expand(SUPPORTED_ARCHITECTURES_WITH_MODEL_ID.items())
+    @require_torch_gpu
+    def test_pipeline_on_gpu(self, *args, **kwargs):
+        model_arch, model_id = args
+        onnx_model = ORTModelForFeatureExtraction.from_pretrained(model_id, from_transformers=True)
+        tokenizer = AutoTokenizer.from_pretrained(model_id)
+        pp = pipeline("feature-extraction", model=onnx_model, tokenizer=tokenizer, device=0)
+        text = "My Name is Philipp and i live in Germany."
+        outputs = pp(text)
+        # check model device
+        self.assertEqual(pp.model.device.type.lower(), "cuda")
+        # compare model output class
+        self.assertTrue(any(any(isinstance(item, float) for item in row) for row in outputs[0]))
+
+    @parameterized.expand(SUPPORTED_ARCHITECTURES_WITH_MODEL_ID.items())
+    def test_default_pipeline_and_model_device(self, *args, **kwargs):
+        model_arch, model_id = args
+        onnx_model = ORTModelForFeatureExtraction.from_pretrained(model_id, from_transformers=True)
+        tokenizer = AutoTokenizer.from_pretrained(model_id)
+        pp = pipeline("feature-extraction", model=onnx_model, tokenizer=tokenizer)
+        self.assertEqual(pp.device, onnx_model.device)
+
 
-class ORTModelForCausalLMIntergrationTest(unittest.TestCase):
+class ORTModelForCausalLMIntegrationTest(unittest.TestCase):
     SUPPORTED_ARCHITECTURES_WITH_MODEL_ID = {
         "gpt2": "hf-internal-testing/tiny-random-gpt2",
         "distilgpt2": "distilgpt2",
@@ -495,3 +600,26 @@ def test_pipeline(self, *args, **kwargs):
         # compare model output class
         self.assertTrue(isinstance(outputs[0]["generated_text"], str))
         self.assertTrue(len(outputs[0]["generated_text"]) > len(text))
+
+    @parameterized.expand(SUPPORTED_ARCHITECTURES_WITH_MODEL_ID.items())
+    @require_torch_gpu
+    def test_pipeline_on_gpu(self, *args, **kwargs):
+        model_arch, model_id = args
+        onnx_model = ORTModelForCausalLM.from_pretrained(model_id, from_transformers=True)
+        tokenizer = AutoTokenizer.from_pretrained(model_id)
+        pp = pipeline("text-generation", model=onnx_model, tokenizer=tokenizer)
+        text = "My Name is Philipp and i live"
+        outputs = pp(text)
+        # check model device
+        self.assertEqual(pp.model.device.type.lower(), "cuda")
+        # compare model output class
+        self.assertTrue(isinstance(outputs[0]["generated_text"], str))
+        self.assertTrue(len(outputs[0]["generated_text"]) > len(text))
+
+    @parameterized.expand(SUPPORTED_ARCHITECTURES_WITH_MODEL_ID.items())
+    def test_default_pipeline_and_model_device(self, *args, **kwargs):
+        model_arch, model_id = args
+        onnx_model = ORTModelForCausalLM.from_pretrained(model_id, from_transformers=True)
+        tokenizer = AutoTokenizer.from_pretrained(model_id)
+        pp = pipeline("text-generation", model=onnx_model, tokenizer=tokenizer)
+        self.assertEqual(pp.device, onnx_model.device)
diff --git a/tests/onnxruntime/test_utils.py b/tests/onnxruntime/test_utils.py
new file mode 100644
--- /dev/null
+++ b/tests/onnxruntime/test_utils.py
@@ -0,0 +1,15 @@
+import unittest
+
+import torch
+
+from optimum.onnxruntime.utils import get_device_for_provider, get_provider_for_device
+
+
+class TestProviderAndDeviceGetters(unittest.TestCase):
+    def test_get_device_for_provider(self):
+        self.assertEqual(get_device_for_provider("CPUExecutionProvider"), torch.device("cpu"))
+        self.assertEqual(get_device_for_provider("CUDAExecutionProvider"), torch.device("cuda"))
+
+    def test_get_provider_for_device(self):
+        self.assertEqual(get_provider_for_device(torch.device("cpu")), "CPUExecutionProvider")
+        self.assertEqual(get_provider_for_device(torch.device("cuda")), "CUDAExecutionProvider")