Repository: pytorch
Issue: [ONNX] torch.onnx.export snapshots the grads as constants in onnx when op is in cuda device
Labels: module: onnx,triaged

### 🐛 Describe the bug

When exporting a module whose forward function includes some `torch.autograd.grad` call, torch.onnx.export saves the results tensor of `torch.autograd.grad` call as constants, instead of saves the nodes which computes the grad to onnx.

This seems to be a bug.

```python
import torch
from torch import nn

class MyModule(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.conv = nn.Linear(16, 15)
    def forward(self, x):
        x = x.detach().requires_grad_(True)
        y = self.conv(x).sum()
        g = torch.autograd.grad(y, x)[0]
        return y-g

def export_model(device):
    model = MyModule().to(device)
    def freeze_model(m):
        for p in m.parameters():
            p.requires_grad_(False)
        return m
    model = freeze_model(model)
    torch.onnx.export(model, torch.ones([1,4,16,16]).to(device), f"linear-{device}.onnx", verbose=True)

export_model('cpu')

export_model('cuda')
```

Exported onnx model of cpu looks like this 
![linear-cpu](https://user-images.githubusercontent.com/8744749/203906301-469434c4-1afb-483f-b556-1e7a0ed9ab99.PNG)

We can see that the backward function of linear is recorded as a GEMM in onnx, which is expected.

Export  onnx model of cuda looks like this 
![linear-cuda](https://user-images.githubusercontent.com/8744749/203906322-e919bac1-0c42-4e3d-8fd7-691fdbf41f65.PNG)

This onnx model is different with the cpu version, it snapshots the grads as a constants tensors, and attached it to the input B of the "Sub" node.





### Versions

Collecting environment information...
PyTorch version: 1.13.0
Is debug build: False
CUDA used to build PyTorch: 11.7
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.5 LTS (x86_64)
GCC version: (Ubuntu 9.4.0-1ubuntu1~20.04.1) 9.4.0
Clang version: 10.0.0-4ubuntu1 
CMake version: version 3.22.5
Libc version: glibc-2.31

Python version: 3.10.8 (main, Nov  4 2022, 13:48:29) [GCC 11.2.0] (64-bit runtime)
Python platform: Linux-5.4.0-131-generic-x86_64-with-glibc2.31
Is CUDA available: True
CUDA runtime version: 11.7.99
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: 
GPU 0: NVIDIA Graphics Device
GPU 1: Tesla T4
GPU 2: NVIDIA T4 32GB
GPU 3: Tesla V100S-PCIE-32GB
GPU 4: Quadro RTX 8000
GPU 5: Tesla P4

Nvidia driver version: 520.40
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] numpy==1.23.4
[pip3] torch==1.13.0
[pip3] torchaudio==0.13.0
[pip3] torchvision==0.14.0
[conda] blas                      1.0                         mkl  
[conda] ffmpeg                    4.3                  hf484d3e_0    pytorch
[conda] mkl                       2021.4.0           h06a4308_640  
[conda] mkl-service               2.4.0           py310h7f8727e_0  
[conda] mkl_fft                   1.3.1           py310hd6ae3a3_0  
[conda] mkl_random                1.2.2           py310h00e6091_0  
[conda] numpy                     1.23.4          py310hd5efca6_0  
[conda] numpy-base                1.23.4          py310h8e6c178_0  
[conda] pytorch                   1.13.0          py3.10_cuda11.7_cudnn8.5.0_0    pytorch
[conda] pytorch-cuda              11.7                 h67b0de4_0    pytorch
[conda] pytorch-mutex             1.0                        cuda    pytorch
[conda] torchaudio                0.13.0              py310_cu117    pytorch
[conda] torchvision               0.14.0              py310_cu117    pytorch