Repository: huggingface/optimum (Python)

Issue:
`device_id` in `provider_options` is ignored
### System Info

```shell
optimum from source (v1.9.0)
```


### Who can help?

_No response_

### Information

- [ ] The official example scripts
- [ ] My own modified scripts

### Tasks

- [X] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

```python
from optimum.onnxruntime import ORTModelForCausalLM
import torch

model = ORTModelForCausalLM.from_pretrained('gpt2', export=True, provider='CUDAExecutionProvider', provider_options={'device_id': 1})

input_ids = torch.randint(0, 100, (1, 1), dtype=torch.long, device='cuda:1')
attention_mask = torch.ones_like(dummy_input_ids, dtype=torch.long, device='cuda:1')

_ = model(input_ids, attention_mask=attention_mask)
_ = model.generate(input_ids, max_new_tokens=10)
```

raises `RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cuda:1!`

### Expected behavior

```python
from optimum.onnxruntime import ORTModelForCausalLM
import torch

model = ORTModelForCausalLM.from_pretrained('gpt2', export=True, provider='CUDAExecutionProvider', provider_options={'device_id': 1})

model.device = torch.device('cuda:1')

dummy_input_ids = torch.randint(0, 100, (1, 1), dtype=torch.long, device='cuda:1')
dummy_attention_mask = torch.ones_like(dummy_input_ids, dtype=torch.long, device='cuda:1')

_ = model(dummy_input_ids, attention_mask=dummy_attention_mask)
_ = model.generate(dummy_input_ids, max_new_tokens=10)
```
works fine


Discussion: Thanks, `self.device` is always set to `cuda:0` with CUDA EP. I'll fix.



Fix (patch):
diff --git a/optimum/onnxruntime/modeling_ort.py b/optimum/onnxruntime/modeling_ort.py
index 10d6beb118..54c3143cc1 100644
--- a/optimum/onnxruntime/modeling_ort.py
+++ b/optimum/onnxruntime/modeling_ort.py
@@ -207,7 +207,9 @@ def shared_attributes_init(
             )
 
         self.providers = model.get_providers()
-        self._device = get_device_for_provider(self.providers[0])
+        self._device = get_device_for_provider(
+            self.providers[0], provider_options=model.get_provider_options()[self.providers[0]]
+        )
 
         # This attribute is needed to keep one reference on the temporary directory, since garbage collecting it
         # would end-up removing the directory containing the underlying ONNX model.
diff --git a/optimum/onnxruntime/utils.py b/optimum/onnxruntime/utils.py
index 02fbf161f0..f8f6acbbdb 100644
--- a/optimum/onnxruntime/utils.py
+++ b/optimum/onnxruntime/utils.py
@@ -173,15 +173,14 @@ def wrap_onnx_config_for_loss(onnx_config: OnnxConfig) -> OnnxConfig:
     return OnnxConfigWithLoss(onnx_config)
 
 
-def get_device_for_provider(provider: str) -> torch.device:
+def get_device_for_provider(provider: str, provider_options: Dict) -> torch.device:
     """
     Gets the PyTorch device (CPU/CUDA) associated with an ONNX Runtime provider.
     """
-    return (
-        torch.device("cuda:0")
-        if provider in ["CUDAExecutionProvider", "TensorrtExecutionProvider"]
-        else torch.device("cpu")
-    )
+    if provider in ["CUDAExecutionProvider", "TensorrtExecutionProvider"]:
+        return torch.device(f"cuda:{provider_options['device_id']}")
+    else:
+        return torch.device("cpu")
 
 
 def get_provider_for_device(device: torch.device) -> str: