Repository: pytorch
Issue: [LazyTensor] model after to(device), traced IR changed and lose type info
Labels: triaged,lazy,Stale

## 🐛 Bug

a model after `model.to("xla")` and then `model.to("cpu")`, then the traced IR will be different with before that and the type information are erase.
I test with bert and roberta in transformers. but there is no issue for resnet

Seems `model.to(...)` only move params/buffers to different device, why this happens?

## To Reproduce

```python
model = ...
traced = torch.jit.trace(model, inp)
g = _inline(traced.graph.copy())
print(g)
model = model.to("xla")  # hacked for our own backend, it will run on cpu/gpu
model = model.to("cpu")
traced = torch.jit.trace(model, inp)
g = _inline(traced.graph.copy())
print(g)
```

printed 1st:
```
graph(%self.1 : __torch__.mnm.frontend.pytorch.TraceWrapper,
      %inp : Long(64, 128, strides=[128, 1], requires_grad=0, device=cpu)):
  %2 : __torch__.transformers.models.roberta.modeling_roberta.RobertaForMaskedLM = prim::GetAttr[name="model"](%self.1)
  %3 : __torch__.transformers.models.roberta.modeling_roberta.RobertaModel = prim::GetAttr[name="roberta"](%2)
  %4 : __torch__.transformers.models.roberta.modeling_roberta.RobertaEmbeddings = prim::GetAttr[name="embeddings"](%3)
  %5 : __torch__.torch.nn.modules.sparse.Embedding = prim::GetAttr[name="word_embeddings"](%4)
  %6 : Tensor = prim::GetAttr[name="weight"](%5)
  %7 : __torch__.transformers.models.roberta.modeling_roberta.RobertaLMHead = prim::GetAttr[name="lm_head"](%2)
  %8 : __torch__.transformers.models.roberta.modeling_roberta.RobertaModel = prim::GetAttr[name="roberta"](%2)
  %9 : __torch__.transformers.models.roberta.modeling_roberta.RobertaEncoder = prim::GetAttr[name="encoder"](%8)
  %10 : __torch__.transformers.models.roberta.modeling_roberta.RobertaEmbeddings = prim::GetAttr[name="embeddings"](%8)
  %11 : int = prim::Constant[value=0](), scope: __module.model/__module.model.roberta # .../python3.7/site-packages/transformers/models/roberta/modeling_roberta.py:760:0
  %12 : int = aten::size(%inp, %11), scope: __module.model/__module.model.roberta # .../python3.7/site-packages/transformers/models/roberta/modeling_roberta.py:760:0
  %batch_size : Long(device=cpu) = prim::NumToTensor(%12), scope: __module.model/__module.model.roberta
  %14 : int = aten::Int(%batch_size), scope: __module.model/__module.model.roberta
  %15 : int = aten::Int(%batch_size), scope: __module.model/__module.model.roberta
  %16 : int = prim::Constant[value=1](), scope: __module.model/__module.model.roberta # .../python3.7/site-packages/transformers/models/roberta/modeling_roberta.py:760:0
  %17 : int = aten::size(%inp, %16), scope: __module.model/__module.model.roberta # .../python3.7/site-packages/transformers/models/roberta/modeling_roberta.py:760:0
  %seq_length : Long(device=cpu) = prim::NumToTensor(%17), scope: __module.model/__module.model.roberta
  %19 : int = aten::Int(%seq_length), scope: __module.model/__module.model.roberta
  %20 : Long(requires_grad=0, device=cpu) = prim::Constant[value={0}](), scope: __module.model/__module.model.roberta # .../python3.7/site-packages/transformers/models/roberta/modeling_roberta.py:774:0
  %21 : int = prim::Constant[value=1](), scope: __module.model/__module.model.roberta # .../python3.7/site-packages/transformers/models/roberta/modeling_roberta.py:774:0
  %22 : Long(requires_grad=0, device=cpu) = aten::add(%seq_length, %20, %21), scope: __module.model/__module.model.roberta # .../python3.7/site-packages/transformers/models/roberta/modeling_roberta.py:774:0
  %23 : int = aten::Int(%22), scope: __module.model/__module.model.roberta
  %24 : int[] = prim::ListConstruct(%15, %23), scope: __module.model/__module.model.roberta
  %25 : int = prim::Constant[value=6](), scope: __module.model/__module.model.roberta # .../python3.7/site-packages/transformers/models/roberta/modeling_roberta.py:774:0
  %26 : NoneType = prim::Constant(), scope: __module.model/__module.model.roberta
  %27 : Device = prim::Constant[value="cpu"](), scope: __module.model/__module.model.roberta # .../python3.7/site-packages/transformers/models/roberta/modeling_roberta.py:774:0
  %28 : bool = prim::Constant[value=0](), scope: __module.model/__module.model.roberta # .../python3.7/site-packages/transformers/models/roberta/modeling_roberta.py:774:0
  %attention_mask.1 : Float(64, 128, strides=[128, 1], requires_grad=0, device=cpu) = aten::ones(%24, %25, %26, %27, %28), scope: __module.model/__module.model.roberta # .../python3.7/site-packages/transformers/models/roberta/modeling_roberta.py:774:0
  %30 : int[] = prim::ListConstruct(%14, %19), scope: __module.model/__module.model.roberta
  %31 : int = prim::Constant[value=4](), scope: __module.model/__module.model.roberta # .../python3.7/site-packages/transformers/models/roberta/modeling_roberta.py:776:0
  %32 : NoneType = prim::Constant(), scope: __module.model/__module.model.roberta
  %33 : Device = prim::Constant[value="cpu"](), scope: __module.model/__module.model.roberta # .../python3.7/site-packages/transformers/models/roberta/modeling_roberta.py:776:0
  %34 : bool = prim::Constant[value=0](), scope: __module.model/__module.model.roberta # .../python3.7/site-packages/transformers/models/roberta/modeling_roberta.py:776:0
  %input.1 : Long(64, 128, strides=[128, 1], requires_grad=0, device=cpu) = aten::zeros(%30, %31, %32, %33, %34), scope: __module.model/__module.model.roberta # .../python3.7/site-packages/transformers/models/roberta/modeling_roberta.py:776:0
  %36 : int = prim::Constant[value=0](), scope: __module.model/__module.model.roberta # .../python3.7/site-packages/transformers/modeling_utils.py:258:0
  %37 : int = prim::Constant[value=0](), scope: __module.model/__module.model.roberta # .../python3.7/site-packages/transformers/modeling_utils.py:258:0
  %38 : int = prim::Constant[value=9223372036854775807](), scope: __module.model/__module.model.roberta # .../python3.7/site-packages/transformers/modeling_utils.py:258:0
  %39 : int = prim::Constant[value=1](), scope: __module.model/__module.model.roberta # .../python3.7/site-packages/transformers/modeling_utils.py:258:0
  %40 : Float(64, 128, strides=[128, 1], requires_grad=0, device=cpu) = aten::slice(%attention_mask.1, %36, %37, %38, %39), scope: __module.model/__module.model.roberta # .../python3.7/site-packages/transformers/modeling_utils.py:258:0
  %41 : int = prim::Constant[value=1](), scope: __module.model/__module.model.roberta # .../python3.7/site-packages/transformers/modeling_utils.py:258:0
  %42 : Float(64, 1, 128, strides=[128, 128, 1], requires_grad=0, device=cpu) = aten::unsqueeze(%40, %41), scope: __module.model/__module.model.roberta # .../python3.7/site-packages/transformers/modeling_utils.py:258:0
  %43 : int = prim::Constant[value=2](), scope: __module.model/__module.model.roberta # .../python3.7/site-packages/transformers/modeling_utils.py:258:0
  %44 : Float(64, 1, 1, 128, strides=[128, 128, 128, 1], requires_grad=0, device=cpu) = aten::unsqueeze(%42, %43), scope: __module.model/__module.model.roberta # .../python3.7/site-packages/transformers/modeling_utils.py:258:0
  %45 : int = prim::Constant[value=3](), scope: __module.model/__module.model.roberta # .../python3.7/site-packages/transformers/modeling_utils.py:258:0
  %46 : int = prim::Constant[value=0](), scope: __module.model/__module.model.roberta # .../python3.7/site-packages/transformers/modeling_utils.py:258:0
  %47 : int = prim::Constant[value=9223372036854775807](), scope: __module.model/__module.model.roberta # .../python3.7/site-packages/transformers/modeling_utils.py:258:0
  %48 : int = prim::Constant[value=1](), scope: __module.model/__module.model.roberta # .../python3.7/site-packages/transformers/modeling_utils.py:258:0
  %extended_attention_mask : Float(64, 1, 1, 128, strides=[128, 128, 128, 1], requires_grad=0, device=cpu) = aten::slice(%44, %45, %46, %47, %48), scope: __module.model/__module.model.roberta # .../python3.7/site-packages/transformers/modeling_utils.py:258:0
  %50 : int = prim::Constant[value=6](), scope: __module.model/__module.model.roberta # .../python3.7/site-packages/transformers/modeling_utils.py:271:0
  %51 : bool = prim::Constant[value=0](), scope: __module.model/__module.model.roberta # .../python3.7/site-packages/transformers/modeling_utils.py:271:0
  %52 : bool = prim::Constant[value=0](), scope: __module.model/__module.model.roberta # .../python3.7/site-packages/transformers/modeling_utils.py:271:0
  %53 : NoneType = prim::Constant(), scope: __module.model/__module.model.roberta
  %54 : Float(64, 1, 1, 128, strides=[128, 128, 128, 1], requires_grad=0, device=cpu) = aten::to(%extended_attention_mask, %50, %51, %52, %53), scope: __module.model/__module.model.roberta # .../python3.7/site-packages/transformers/modeling_utils.py:271:0
  %55 : float = prim::Constant[value=1.](), scope: __module.model/__module.model.roberta # .../python3.7/site-packages/torch/_tensor.py:544:0
  %56 : int = prim::Constant[value=1](), scope: __module.model/__module.model.roberta # .../python3.7/site-packages/torch/_tensor.py:544:0
  %57 : Float(64, 1, 1, 128, strides=[128, 128, 128, 1], requires_grad=0, device=cpu) = aten::rsub(%54, %55, %56), scope: __module.model/__module.model.roberta # .../python3.7/site-packages/torch/_tensor.py:544:0
  %58 : Double(requires_grad=0, device=cpu) = prim::Constant[value={-10000}](), scope: __module.model/__module.model.roberta # .../python3.7/site-packages/transformers/modeling_utils.py:272:0
  %attention_mask : Float(64, 1, 1, 128, strides=[128, 128, 128, 1], requires_grad=0, device=cpu) = aten::mul(%57, %58), scope: __module.model/__module.model.roberta # .../python3.7/site-packages/transformers/modeling_utils.py:272:0
  %60 : __torch__.torch.nn.modules.dropout.Dropout = prim::GetAttr[name="dropout"](%10)
  %61 : __torch__.torch.nn.modules.normalization.LayerNorm = prim::GetAttr[name="LayerNorm"](%10)
  %62 : __torch__.torch.nn.modules.sparse.___torch_mangle_0.Embedding = prim::GetAttr[name="position_embeddings"](%10)
  %63 : __torch__.torch.nn.modules.sparse.___torch_mangle_1.Embedding = prim::GetAttr[name="token_type_embeddings"](%10)
  %64 : __torch__.torch.nn.modules.sparse.Embedding = prim::GetAttr[name="word_embeddings"](%10)
  %65 : int = prim::Constant[value=1](), scope: __module.model/__module.model.roberta/__module.model.roberta.embeddings # .../python3.7/site-packages/transformers/models/roberta/modeling_roberta.py:1497:0
  %66 : Bool(64, 128, strides=[128, 1], requires_grad=0, device=cpu) = aten::ne(%inp, %65), scope: __module.model/__module.model.roberta/__module.model.roberta.embeddings # .../python3.7/site-packages/transformers/models/roberta/modeling_roberta.py:1497:
...
...
```

2nd:

```
graph(%self.1 : __torch__.mnm.frontend.pytorch.TraceWrapper,
      %inp.1 : Tensor):
  %2 : __torch__.transformers.models.roberta.modeling_roberta.RobertaForMaskedLM = prim::GetAttr[name="model"](%self.1)
  %3 : __torch__.transformers.models.roberta.modeling_roberta.RobertaLMHead = prim::GetAttr[name="lm_head"](%2)
  %4 : __torch__.transformers.models.roberta.modeling_roberta.RobertaModel = prim::GetAttr[name="roberta"](%2)
  %5 : Tensor = prim::Constant[value={-10000}]() # .../python3.7/site-packages/transformers/modeling_utils.py:272:0
  %6 : bool = prim::Constant[value=0]() # .../python3.7/site-packages/transformers/models/roberta/modeling_roberta.py:774:0
  %7 : Device = prim::Constant[value="cpu"]()
  %8 : NoneType = prim::Constant() # :0:0
  %9 : Tensor = prim::Constant[value={0}]() # .../python3.7/site-packages/transformers/models/roberta/modeling_roberta.py:774:0
  %10 : int = prim::Constant[value=0]() # .../python3.7/site-packages/transformers/models/roberta/modeling_roberta.py:760:0
  %11 : int = prim::Constant[value=1]() # .../python3.7/site-packages/transformers/models/roberta/modeling_roberta.py:760:0
  %12 : int = prim::Constant[value=6]() # .../python3.7/site-packages/transformers/models/roberta/modeling_roberta.py:774:0
  %13 : int = prim::Constant[value=4]() # .../python3.7/site-packages/transformers/models/roberta/modeling_roberta.py:776:0
  %14 : int = prim::Constant[value=9223372036854775807]() # .../python3.7/site-packages/transformers/modeling_utils.py:258:0
  %15 : int = prim::Constant[value=2]() # .../python3.7/site-packages/transformers/modeling_utils.py:258:0
  %16 : int = prim::Constant[value=3]() # .../python3.7/site-packages/transformers/modeling_utils.py:258:0
  %17 : float = prim::Constant[value=1.]() # .../python3.7/site-packages/torch/_tensor.py:544:0
  %18 : __torch__.transformers.models.roberta.modeling_roberta.RobertaEncoder = prim::GetAttr[name="encoder"](%4)
  %19 : __torch__.transformers.models.roberta.modeling_roberta.RobertaEmbeddings = prim::GetAttr[name="embeddings"](%4)
  %20 : int = aten::size(%inp.1, %10) # .../python3.7/site-packages/transformers/models/roberta/modeling_roberta.py:760:0
  %batch_size.1 : Tensor = prim::NumToTensor(%20) # :0:0
  %22 : int = aten::Int(%batch_size.1)
  %23 : int = aten::Int(%batch_size.1)
  %24 : int = aten::size(%inp.1, %11) # .../python3.7/site-packages/transformers/models/roberta/modeling_roberta.py:760:0
  %seq_length.1 : Tensor = prim::NumToTensor(%24) # :0:0
  %26 : int = aten::Int(%seq_length.1)
  %27 : Tensor = aten::add(%seq_length.1, %9, %11) # .../python3.7/site-packages/transformers/models/roberta/modeling_roberta.py:774:0
  %28 : int = aten::Int(%27)
  %29 : int[] = prim::ListConstruct(%23, %28)
  %attention_mask.1 : Tensor = aten::ones(%29, %12, %8, %7, %6) # .../python3.7/site-packages/transformers/models/roberta/modeling_roberta.py:774:0
  %31 : int[] = prim::ListConstruct(%22, %26)
  %input.14 : Tensor = aten::zeros(%31, %13, %8, %7, %6) # .../python3.7/site-packages/transformers/models/roberta/modeling_roberta.py:776:0
  %33 : Tensor = aten::slice(%attention_mask.1, %10, %10, %14, %11) # .../python3.7/site-packages/transformers/modeling_utils.py:258:0
  %34 : Tensor = aten::unsqueeze(%33, %11) # .../python3.7/site-packages/transformers/modeling_utils.py:258:0
  %35 : Tensor = aten::unsqueeze(%34, %15) # .../python3.7/site-packages/transformers/modeling_utils.py:258:0
  %extended_attention_mask.1 : Tensor = aten::slice(%35, %16, %10, %14, %11) # .../python3.7/site-packages/transformers/modeling_utils.py:258:0
  %37 : Tensor = aten::to(%extended_attention_mask.1, %12, %6, %6, %8) # .../python3.7/site-packages/transformers/modeling_utils.py:271:0
  %38 : Tensor = aten::rsub(%37, %17, %11) # .../python3.7/site-packages/torch/_tensor.py:544:0
  %attention_mask0.1 : Tensor = aten::mul(%38, %5) # .../python3.7/site-packages/transformers/modeling_utils.py:272:0
  %40 : Device = prim::Constant[value="cpu"]()
  %41 : Tensor = prim::Constant[value={1}]() # .../python3.7/site-packages/transformers/models/roberta/modeling_roberta.py:1499:0
  %42 : Tensor = prim::Constant[value={0}]() # .../python3.7/site-packages/transformers/models/roberta/modeling_roberta.py:1498:0
  %43 : NoneType = prim::Constant()
  %44 : bool = prim::Constant[value=0]()
  %45 : int = prim::Constant[value=1]() # .../python3.7/site-packages/transformers/models/roberta/modeling_roberta.py:1497:0
  %46 : int = prim::Constant[value=3]() # .../python3.7/site-packages/transformers/models/roberta/modeling_roberta.py:1497:0
  %47 : int = prim::Constant[value=4]() # .../python3.7/site-packages/transformers/models/roberta/modeling_roberta.py:1499:0
  %48 : int = prim::Constant[value=0]() # .../python3.7/site-packages/transformers/models/roberta/modeling_roberta.py:103:0
  %49 : __torch__.torch.nn.modules.dropout.Dropout = prim::GetAttr[name="dropout"](%19)
  %50 : __torch__.torch.nn.modules.normalization.LayerNorm = prim::GetAttr[name="LayerNorm"](%19)
  %51 : __torch__.torch.nn.modules.sparse.___torch_mangle_0.Embedding = prim::GetAttr[name="position_embeddings"](%19)
  %52 : __torch__.torch.nn.modules.sparse.___torch_mangle_1.Embedding = prim::GetAttr[name="token_type_embeddings"](%19)
  %53 : __torch__.torch.nn.modules.sparse.Embedding = prim::GetAttr[name="word_embeddings"](%19)
  %54 : Tensor = aten::ne(%inp.1, %45) # .../python3.7/site-packages/transformers/models/roberta/modeling_roberta.py:1497:0
...
...
```

There is no type annotation and the number of constant changes.

## Environment
torch is on branch `lazy_tensor_staging`, commit id is `0e8776b4d45a243df6e8499d070e2df89dcad1f9`, build with out CUDA

PyTorch version: 1.9.0a0+git0e8776b
Is debug build: False
CUDA used to build PyTorch: Could not collect
ROCM used to build PyTorch: N/A

OS: Ubuntu 18.04.6 LTS (x86_64)
GCC version: (Ubuntu 7.5.0-3ubuntu1~18.04) 7.5.0
Clang version: Could not collect
CMake version: version 3.21.4
Libc version: glibc-2.17

Python version: 3.7.11 (default, Jul 27 2021, 14:32:16)  [GCC 7.5.0] (64-bit runtime)
Python platform: Linux-5.4.0-1059-aws-x86_64-with-debian-buster-sid
Is CUDA available: False
CUDA runtime version: 10.2.89
GPU models and configuration: GPU 0: Tesla V100-SXM2-16GB
Nvidia driver version: 450.142.00
cuDNN version: Probably one of the following:
/usr/local/cuda-10.1/targets/x86_64-linux/lib/libcudnn.so.7.6.5
/usr/local/cuda-10.2/targets/x86_64-linux/lib/libcudnn.so.7.6.5
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn.so.8.0.5
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn_adv_infer.so.8.0.5
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn_adv_train.so.8.0.5
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn_cnn_infer.so.8.0.5
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn_cnn_train.so.8.0.5
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn_ops_infer.so.8.0.5
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn_ops_train.so.8.0.5
/usr/local/cuda-11.1/targets/x86_64-linux/lib/libcudnn.so.8.0.5
/usr/local/cuda-11.1/targets/x86_64-linux/lib/libcudnn_adv_infer.so.8.0.5
/usr/local/cuda-11.1/targets/x86_64-linux/lib/libcudnn_adv_train.so.8.0.5
/usr/local/cuda-11.1/targets/x86_64-linux/lib/libcudnn_cnn_infer.so.8.0.5
/usr/local/cuda-11.1/targets/x86_64-linux/lib/libcudnn_cnn_train.so.8.0.5
/usr/local/cuda-11.1/targets/x86_64-linux/lib/libcudnn_ops_infer.so.8.0.5
/usr/local/cuda-11.1/targets/x86_64-linux/lib/libcudnn_ops_train.so.8.0.5
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] numpy==1.21.4
[pip3] torch==1.9.0a0+git0e8776b
[pip3] torch-mnm==1.9
[pip3] torchvision==0.11.1
[conda] numpy                     1.21.4                   pypi_0    pypi
[conda] torch                     1.9.0a0+git0e8776b          pypi_0    pypi
[conda] torch-mnm                 1.9                      pypi_0    pypi
[conda] torchvision               0.11.1                   pypi_0    pypi