Repository: pytorch
Issue: Strange behavior of torch.jit.trace when moving parameters across device
Labels: oncall: jit,Stale

### 🐛 Describe the bug

Hi, I want to lazily load the model parameters and trace the model.

However, it seems that it should be moved **twice** to work properly with `torch.jit.trace`. (Uncomment the code block and it will work)

Additionally, I notice that it consumes a lot of memory (cpu) with lazily load parameter + trace. Am I using this tool incorrectly? And what is a good practice to do this? 

Any feedback would be greatly appreciated! :)

```python
import torch

def release_weights(module, *args, **kwargs):
    module.to('cpu')

def load_weights(module, *args, **kwargs):
    module.to('cuda')

model = torch.nn.Linear(1, 1).eval()
model.register_forward_pre_hook(load_weights)
model.register_forward_hook(release_weights)

# Uncomment to add hooks twice, then it should work
# model.register_forward_pre_hook(load_weights)
# model.register_forward_hook(release_weights)

data = torch.ones(1, device='cuda')
model = torch.jit.trace(model, data)

model(data)
```

```text
Traceback (most recent call last):
  File "/home/siahuat0727/pylomin/env/lib/python3.9/site-packages/torch/jit/_trace.py", line 443, in run_mod_and_filter_tensor_outputs
    outs = wrap_retval(mod(*_clone_inputs(inputs)))
RuntimeError: The following operation failed in the TorchScript interpreter.
Traceback of TorchScript (most recent call last):
/home/siahuat0727/pylomin/env/lib/python3.9/site-packages/torch/nn/functional.py(1848): linear
/home/siahuat0727/pylomin/env/lib/python3.9/site-packages/torch/nn/modules/linear.py(103): forward
/home/siahuat0727/pylomin/env/lib/python3.9/site-packages/torch/nn/modules/module.py(1090): _slow_forward
/home/siahuat0727/pylomin/env/lib/python3.9/site-packages/torch/nn/modules/module.py(1120): _call_impl
/home/siahuat0727/pylomin/env/lib/python3.9/site-packages/torch/jit/_trace.py(958): trace_module
/home/siahuat0727/pylomin/env/lib/python3.9/site-packages/torch/jit/_trace.py(741): trace
/home/siahuat0727/pylomin/test_hook.py(21): <module>
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 mat2 in method wrapper_mm)


The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/home/siahuat0727/pylomin/test_hook.py", line 21, in <module>
    model = torch.jit.trace(model, data)
  File "/home/siahuat0727/pylomin/env/lib/python3.9/site-packages/torch/jit/_trace.py", line 741, in trace
    return trace_module(
  File "/home/siahuat0727/pylomin/env/lib/python3.9/site-packages/torch/jit/_trace.py", line 983, in trace_module
    _check_trace(
  File "/home/siahuat0727/pylomin/env/lib/python3.9/site-packages/torch/autograd/grad_mode.py", line 28, in decorate_context
    return func(*args, **kwargs)
  File "/home/siahuat0727/pylomin/env/lib/python3.9/site-packages/torch/jit/_trace.py", line 516, in _check_trace
    traced_outs = run_mod_and_filter_tensor_outputs(traced_func, inputs, "trace")
  File "/home/siahuat0727/pylomin/env/lib/python3.9/site-packages/torch/jit/_trace.py", line 449, in run_mod_and_filter_tensor_outputs
    raise TracingCheckError(
torch.jit._trace.TracingCheckError: Tracing failed sanity checks!
encountered an exception while running the trace with test inputs.
Exception:
        The following operation failed in the TorchScript interpreter.
        Traceback of TorchScript (most recent call last):
        /home/siahuat0727/pylomin/env/lib/python3.9/site-packages/torch/nn/functional.py(1848): linear
        /home/siahuat0727/pylomin/env/lib/python3.9/site-packages/torch/nn/modules/linear.py(103): forward
        /home/siahuat0727/pylomin/env/lib/python3.9/site-packages/torch/nn/modules/module.py(1090): _slow_forward
        /home/siahuat0727/pylomin/env/lib/python3.9/site-packages/torch/nn/modules/module.py(1120): _call_impl
        /home/siahuat0727/pylomin/env/lib/python3.9/site-packages/torch/jit/_trace.py(958): trace_module
        /home/siahuat0727/pylomin/env/lib/python3.9/site-packages/torch/jit/_trace.py(741): trace
        /home/siahuat0727/pylomin/test_hook.py(21): <module>
        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 mat2 in method wrapper_mm)
```

### Versions

```
PyTorch version: 1.10.1+cu102
Is debug build: False
CUDA used to build PyTorch: 10.2
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.15.5
Libc version: glibc-2.27

Python version: 3.9.9 (main, Nov 16 2021, 03:05:18)  [GCC 7.5.0] (64-bit runtime)
Python platform: Linux-5.4.0-1064-azure-x86_64-with-glibc2.27
Is CUDA available: True
CUDA runtime version: 11.1.105
GPU models and configuration: GPU 0: Tesla P100-PCIE-16GB
Nvidia driver version: 470.82.01
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.8.3.1
/usr/lib/x86_64-linux-gnu/libcudnn_adv_infer.so.8.3.1
/usr/lib/x86_64-linux-gnu/libcudnn_adv_train.so.8.3.1
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_infer.so.8.3.1
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_train.so.8.3.1
/usr/lib/x86_64-linux-gnu/libcudnn_ops_infer.so.8.3.1
/usr/lib/x86_64-linux-gnu/libcudnn_ops_train.so.8.3.1
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] numpy==1.21.5
[pip3] torch==1.10.1
[pip3] torchvision==0.11.2
[conda] blas                      1.0                         mkl
[conda] cudatoolkit               11.1.1               h6406543_8    conda-forge
[conda] ffmpeg                    4.3                  hf484d3e_0    pytorch
[conda] mkl                       2021.3.0           h06a4308_520
[conda] mkl-service               2.4.0            py38h7f8727e_0
[conda] mkl_fft                   1.3.0            py38h42c9631_2
[conda] mkl_random                1.2.2            py38h51133e4_0
[conda] mypy-extensions           0.4.3                    pypi_0    pypi
[conda] numpy                     1.21.2                   pypi_0    pypi
[conda] torch                     1.10.0                   pypi_0    pypi
[conda] torch-tb-profiler         0.2.1                    pypi_0    pypi
[conda] torchaudio                0.9.0                      py38    pytorch
[conda] torchtext                 0.10.0                     py38    pytorch
[conda] torchvision               0.10.0               py38_cu111    pytorch
```