Repository: huggingface/optimum

Issue:
torch_ones in fx.optimisation.FuseBiasInLinear creates tensor/proxy in cpu
### System Info

```shell
Colab Environment:

- `optimum` version : 1.5.2
- `transformers` version: 4.25.1
- Platform: Linux-5.10.133+-x86_64-with-glibc2.27
- Python version: 3.8.16
- Huggingface_hub version: 0.11.1
- PyTorch version (GPU?): 1.13.0+cu116 (True)
- Using GPU in script?: True
```


### Who can help?

@regisss 

In this line [here](https://github.com/huggingface/optimum/blob/main/optimum/fx/optimization/transformations.py#L409), device could be specified:

```python
return torch.cat([linear_input, torch_ones(shape, device=linear_input.device)], dim=-1)
```

I tested it.

### Information

- [ ] The official example scripts
- [ ] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

To reproduce this error:

```python
import torch
from transformers import BertModel
from transformers.utils.fx import symbolic_trace

# original Model
original_bert = BertModel.from_pretrained("bert-base-uncased")
original_bert.eval()

# traced Model
traced_bert = symbolic_trace(
    original_bert,
    input_names=["input_ids", "attention_mask", "token_type_ids"],
)

# put models on cuda
original_bert.to(device)
traced_bert.to(device)

# input configuration
device = 'cuda'
bert_inputs = dict()
batch_size, seq_length = 4, 512

# create inputs
bert_inputs['input_ids'] = torch.zeros(batch_size, seq_length, dtype=torch.int, device=device).random_(original_bert.config.vocab_size)
bert_inputs['token_type_ids'] = torch.zeros(batch_size, seq_length, dtype=torch.int, device=device)
bert_inputs['attention_mask'] = torch.ones(batch_size, seq_length, dtype=torch.int, device=device)

# transform graph
transformation = FuseBiasInLinear()
transformed_bert = transformation(traced_bert)
transformed_outputs = transformed_bert(**bert_inputs)
```

## TraceBack:

<details>

  ```python
  Traceback (most recent call last):
    File "/usr/local/lib/python3.8/dist-packages/torch/fx/graph_module.py", line 267, in __call__
      return super(self.cls, obj).__call__(*args, **kwargs)  # type: ignore[misc]
    File "/usr/local/lib/python3.8/dist-packages/torch/nn/modules/module.py", line 1190, in _call_impl
      return forward_call(*input, **kwargs)
    File "<eval_with_key>.39", line 34, in forward
      cat = torch.cat([embeddings_dropout, ones], dim = -1);  ones = None
  RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu! (when checking argument for argument tensors in method wrapper_cat)
  
  Call using an FX-traced Module, line 34 of the traced Module's generated forward function:
      ones = torch.ones(add_87);  add_87 = None
      cat = torch.cat([embeddings_dropout, ones], dim = -1);  ones = None
  
  ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
      encoder_layer_0_attention_self_query = getattr(self.encoder.layer, "0").attention.self.query(cat);  cat = None
  
      getattr_2 = embeddings_dropout.shape
  
  ---------------------------------------------------------------------------
  RuntimeError                              Traceback (most recent call last)
  [<ipython-input-76-19b42e791ede>](https://localhost:8080/#) in <module>
        4 transformed_bert = composition(traced_bert)
        5 transformed_bert.to(device)
  ----> 6 transformed_outputs = transformed_bert(**bert_inputs)
        7 
        8 # verify outputs
  
  1 frames
  [/usr/local/lib/python3.8/dist-packages/torch/fx/graph_module.py](https://localhost:8080/#) in __call__(self, obj, *args, **kwargs)
      273                 print(_WrappedCall._generate_error_message(topmost_framesummary),
      274                       file=sys.stderr)
  --> 275                 raise e.with_traceback(None)
      276             else:
      277                 raise e
  
  RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu! (when checking argument for argument tensors in method wrapper_cat)
  ```

</details>

### Expected behavior

To run without error.


Discussion:
Hi Ilyas, good catch! I'm able to reproduce this error. Would you like to open a PR to fix this?
Yes ofc I would love to!

Resolving patch:
diff --git a/optimum/fx/optimization/transformations.py b/optimum/fx/optimization/transformations.py
--- a/optimum/fx/optimization/transformations.py
+++ b/optimum/fx/optimization/transformations.py
@@ -406,7 +406,7 @@ def transform(self, graph_module: "GraphModule") -> "GraphModule":
 
         def insert_concat(linear_input):
             shape = linear_input.shape[:-1] + (1,)
-            return torch.cat([linear_input, torch_ones(shape)], dim=-1)
+            return torch.cat([linear_input, torch_ones(shape, device=linear_input.device)], dim=-1)
 
         tracer = torch.fx.proxy.GraphAppendingTracer(graph_module.graph)
         for node in graph_module.graph.nodes:


Test patch: