[`pipeline`] Fix str device issue
Author: younesbelkada
Pull request: yes

# What does this PR do?

Addresses: https://github.com/huggingface/transformers/pull/24140#issuecomment-1584617146 

Currently passing `device="cuda"` is not supported when creating a pipeline. 
This is because `torch.cuda.set_device(self.device)` expects the device to have an explicit index. The fix is to create an indexed device when initializing a pipeline with a str device

Handy reproducible snippet:

```python
from transformers import pipeline

# this works
pipe = pipeline("text-generation", device=0)
pipe("Hello")

# this works
pipe = pipeline("text-generation", device="cuda:0")
pipe("Hello")

# this fails
pipe = pipeline("text-generation", device="cuda")
pipe("Hello")
```

cc @amyeroberts @Narsil  

Comments:
_The documentation is not available anymore as the PR was closed or merged._

Also 

```
python -c 'from transformers import pipeline; pipe = pipeline(model="gpt2", device="cuda")'
```

Works on `main`.. So I'm not sure what's the issue


@Narsil what you shared works on main but it should throw an error if you try to run an example with it (I attached a reproducible snippet above)

Alternatively, this fails on main and this PR fixes it

```bash
python -c 'from transformers import pipeline; pipe = pipeline(model="gpt2", device="cuda"); pipe("hello")'
```

Can we remove the `set_device` instead then ? Seems better:

```patch
diff --git a/src/transformers/pipelines/base.py b/src/transformers/pipelines/base.py
index 510c07cf5..b5975d081 100644
--- a/src/transformers/pipelines/base.py
+++ b/src/transformers/pipelines/base.py
@@ -901,10 +901,8 @@ class Pipeline(_ScikitCompat):
             with tf.device("/CPU:0" if self.device == -1 else f"/device:GPU:{self.device}"):
                 yield
         else:
-            if self.device.type == "cuda":
-                torch.cuda.set_device(self.device)
-
-            yield
+            with torch.cuda.device(self.device):
+                yield
```

The initial thing fails indeed, and seems to be linked to the fact that there are multiple `set_device` happening causing issues.

By removing it the issue is indeed removed (but the test you added in the test suite isn't failing on main, and since this is what supposed to catch the regression, this is what I tried :) )


I am happy to revert some of the changes I proposed and add yours, it looks much better. However I have few questions 
1- is it ok to call that context manager if `self.device` is CPU? I think we need a check on top of that to make sure we're not on CPU (similarly as what we had before)

```python
import torch
device = torch.device("cpu")

with torch.cuda.device(device):
    print(torch.randn(1))
```
Throws:
```bash
    raise ValueError('Expected a cuda device, but got: {}'.format(device))
ValueError: Expected a cuda device, but got: cpu
```
EDIT: just `with torch.device(self.device)` seems to work

2- I am not sure but I think the `with device` context manager is only available since PT2.0 no?


> 2- I am not sure but I think the with device context manager is only available since PT2.0 no?

I don't know, all those are very good questions for which I don't have the answer to. I just know that now `set_device` is strongly discouraged so it's probably the source of our issues.

Thanks ! 
I can confirm the context manager doesn't work for PT==1.9 which is [should be supported by us](https://github.com/huggingface/transformers/blob/4c6e42958951ca66a6b498b1afce8d8ad4ac2274/setup.py#L178): 

```python
Traceback (most recent call last):
  File "scratch.py", line 203, in <module>
    with torch.device(device):
AttributeError: __enter__
```

Therefore I just added some changes to ensure backward compatibility with older PT versions. WDYT?

Hi @Narsil 
Let me know if the changes look all good to you, happy to address any additional comments you have 

May I attempt a different thing ?

I think the fix is correct, but I'm wondering if simply relying on `torch.cuda.device` context manager couldn't help remove the need for the compat layer.

Sure yes! 

Cannot push

```patch
diff --git a/src/transformers/pipelines/base.py b/src/transformers/pipelines/base.py
index 626d33a3d..ee117e62a 100644
--- a/src/transformers/pipelines/base.py
+++ b/src/transformers/pipelines/base.py
@@ -50,7 +50,6 @@ if is_torch_available():
     from torch.utils.data import DataLoader, Dataset

     from ..models.auto.modeling_auto import AutoModel
-    from ..pytorch_utils import is_torch_greater_or_equal_than_2_0

     # Re-export for backward compatibility
     from .pt_utils import KeyDataset
@@ -794,16 +793,11 @@ class Pipeline(_ScikitCompat):
             if isinstance(device, torch.device):
                 self.device = device
             elif isinstance(device, str):
-                if device == "cuda" and not is_torch_greater_or_equal_than_2_0:
-                    # for backward compatiblity if using `set_device` and `cuda`
-                    device = f"cuda:{torch.cuda.current_device()}"
                 self.device = torch.device(device)
             elif device < 0:
                 self.device = torch.device("cpu")
-            elif isinstance(device, int):
-                self.device = torch.device(f"cuda:{device}")
             else:
-                raise ValueError(f"Device type not supported. Got {device}")
+                self.device = torch.device(f"cuda:{device}")
         else:
             self.device = device if device is not None else -1
         self.torch_dtype = torch_dtype
@@ -908,13 +902,10 @@ class Pipeline(_ScikitCompat):
             with tf.device("/CPU:0" if self.device == -1 else f"/device:GPU:{self.device}"):
                 yield
         else:
-            if is_torch_greater_or_equal_than_2_0:
-                with torch.device(self.device):
+            if self.device.type == "cuda":
+                with torch.cuda.device(self.device):
                     yield
-            # for backward compatibility
             else:
-                if self.device.type == "cuda":
-                    torch.cuda.set_device(self.device)
                 yield
```

`torch.cuda.device` is defined for torch==1.9 so it should work.

And `torch.device("cpu")` ... well it's the default there's no need to context manage it.

Hi @Narsil 
I am not sure if `with torch.cuda.device(self.device):` is supported for torch<2.0 

https://pytorch.org/tutorials/recipes/recipes/changing_default_device.html

Maybe we should merge this PR for now to unblock also @thomasw21 & @NouamaneTazi . what do you think?

I don't think we're blocked by this. 

> And torch.device("cpu") ... well it's the default there's no need to context manage it.

Not sure of the context of this sentence, but we're overriding the default to `cuda`, so having a context manager to switch back to `cpu` makes sense to me.


https://pytorch.org/docs/1.9.0/generated/torch.cuda.device.html?highlight=torch%20cuda%20device#torch.cuda.device

It is supported from 1.9.0+, at least in the docs.

Great ! agreed with those changes