Repository: pytorch
Issue: [torch.fx] Deserialization Error - TypeError: ones() received an invalid combination of arguments - got (tuple, device=Attribute) 
Labels: high priority,triage review,module: fx,oncall: fx

### 🐛 Describe the bug

Dear Pytorch Community,

Recently, I'm working on my research project and trying to serialize and deserialize the `torch.fx GraphModule object` of `bloom-560m `model.

I used pickle as the package for serialization and deserialization. Initially, I successfully serialized the GraphModule object using `pickle.dumps()`.

However, when I try to deserialize the pickle object using `pickle.loads()`, it shows the error shown below:
```
Traceback (most recent call last):
  File "/Users/junchenzhao/Dist-CPU-Learn/test_inference.py", line 133, in <module>
    deserialized_modules = [pickle.loads(i) for i in serialized_modules]
  File "/Users/junchenzhao/Dist-CPU-Learn/test_inference.py", line 133, in <listcomp>
    deserialized_modules = [pickle.loads(i) for i in serialized_modules]
  File "/Users/junchenzhao/anaconda3/envs/onnx-training/lib/python3.10/site-packages/torch/fx/graph_module.py", line 109, in reduce_graph_module
    return _deserialize_graph_module(forward, body)
  File "/Users/junchenzhao/anaconda3/envs/onnx-training/lib/python3.10/site-packages/torch/fx/graph_module.py", line 168, in _deserialize_graph_module
    graph = KeepModules().trace(com, **tracer_extras)
  File "/Users/junchenzhao/anaconda3/envs/onnx-training/lib/python3.10/site-packages/torch/fx/_symbolic_trace.py", line 778, in trace
    (self.create_arg(fn(*args)),),
  File "<eval_with_key>.7", line 18, in forward
TypeError: ones() received an invalid combination of arguments - got (tuple, device=Attribute), but expected one of:
 * (tuple of ints size, *, tuple of names names, torch.dtype dtype, torch.layout layout, torch.device device, bool pin_memory, bool requires_grad)
 * (tuple of ints size, *, Tensor out, torch.dtype dtype, torch.layout layout, torch.device device, bool pin_memory, bool requires_grad)

```

Part of the GraphModule forward code to be serialized can be seen below, where error happens at the `ones`:
```
def forward(self, input_ids : torch.Tensor):
    1 size = input_ids.size()
    2 getitem = size[0]
    3 getitem_1 = size[1];  size = None
    4 transformer_word_embeddings = self.transformer.word_embeddings(input_ids)
    5 transformer_word_embeddings_layernorm = self.transformer.word_embeddings_layernorm(transformer_word_embeddings);  transformer_word_embeddings = None
    6 getattr_1 = transformer_word_embeddings_layernorm.device
 ----> 7 ones = torch.ones((getitem, getitem_1), device = getattr_1);  getattr_1 = None
    8 size_1 = ones.size()
    9 getitem_2 = size_1[0]
    10 getitem_3 = size_1[1];  size_1 = None
    11 getattr_2 = ones.device
    12 tensor = torch.tensor(0.7071067811865476, device = getattr_2, dtype = torch.float32);  getattr_2 = None
    13 getattr_3 = ones.device
    14 arange = torch.arange(1, 17, device = getattr_3, dtype = torch.int32);  getattr_3 = None
    15 pow_1 = torch.pow(tensor, arange);  tensor = arange = None
    16 cumsum = ones.cumsum(dim = -1)
    17 sub = cumsum - 1;  cumsum = None
    18 mul = sub * ones;  sub = None
    ....

```


### Reproduce
A similar issue can also be reproduced using the code shown below:
```
import torch
from torch.fx import symbolic_trace

def test(input_ids):
    size = input_ids.size()
    getitem = size[0]
    getitem_1 = size[1];
    getattr_1 = input_ids.device
    ones = torch.ones((getitem, getitem_1), device=getattr_1)
    return ones

traced = symbolic_trace(test)
```

Same Error Message as you can see below by running the reproduce coding shown above:
```
      7     getitem_1 = size[1];
      8     getattr_1 = input_ids.device
----> 9     ones = torch.ones((getitem, getitem_1), device=getattr_1)
     10     return ones
     11 

TypeError: ones() received an invalid combination of arguments - got (tuple, device=Attribute), but expected one of:
 * (tuple of ints size, *, tuple of names names, torch.dtype dtype, torch.layout layout, torch.device device, bool pin_memory, bool requires_grad)
 * (tuple of ints size, *, Tensor out, torch.dtype dtype, torch.layout layout, torch.device device, bool pin_memory, bool requires_grad)
```

If it's possible, could anyone help me with this issue?

Thanks!

### Urgency
This is urgent and hope to be resolved as soon as possible.

### Versions

PyTorch version: 2.0.1
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: macOS 13.3.1 (x86_64)
GCC version: Could not collect
Clang version: 14.0.3 (clang-1403.0.22.14.1)
CMake version: version 3.26.3
Libc version: N/A

Python version: 3.10.11 (main, Apr 20 2023, 13:59:00) [Clang 14.0.6 ] (64-bit runtime)
Python platform: macOS-10.16-x86_64-i386-64bit
Is CUDA available: False
CUDA runtime version: No CUDA
CUDA_MODULE_LOADING set to: N/A
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Intel(R) Core(TM) i7-9750H CPU @ 2.60GHz

Versions of relevant libraries:
[pip3] numpy==1.24.3
[pip3] torch==2.0.1
[pip3] torchaudio==2.0.2
[pip3] torchvision==0.15.2
[conda] blas                      1.0                         mkl  
[conda] ffmpeg                    4.3                  h0a44026_0    pytorch
[conda] mkl                       2023.1.0         h59209a4_43558  
[conda] mkl-service               2.4.0           py310h6c40b1e_1  
[conda] mkl_fft                   1.3.6           py310h3ea8b11_1  
[conda] mkl_random                1.2.2           py310h3ea8b11_1  
[conda] numpy                     1.24.3          py310h827a554_1  
[conda] numpy-base                1.24.3          py310ha186be2_1  
[conda] pytorch                   2.0.1                  py3.10_0    pytorch
[conda] torchaudio                2.0.2                 py310_cpu    pytorch
[conda] torchvision               0.15.2                py310_cpu    pytorch

cc @ezyang @gchanan @zou3519 @SherlockNoMad @EikanWang @jgong5 @wenzhe-nrv @msaroufim @wconstab @bdhirsh @anijain2305