Repository: pytorch
Issue: [ONNX] torch.onnx.export can not export the grad of conv when the op is in CPU  
Labels: module: onnx,triaged

### 🐛 Describe the bug

pytorch onnx exportor can not export the backward  of conv2d op when it's in CPU, but it can export it when it's in CUDA device.


```python
import torch
from torch import nn

class MyModule(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.conv = nn.Conv2d(4, 64, (3, 3), stride=(1,1), padding=(1,1), dilation=(1,1), groups=1)
    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"conv-{device}.onnx", verbose=True)

print(torch.__version__)
# should work w/o error
export_model('cuda')

# has RuntimeError: Found an unsupported argument type in the JIT tracer. File a bug report.
export_model('cpu')

```

The error log is
```
Exported graph: graph(%onnx::Conv_0 : Float(1, 4, 16, 16, strides=[1024, 256, 16, 1], requires_grad=0, device=cuda:0),
      %conv.weight : Float(64, 4, 3, 3, strides=[36, 9, 3, 1], requires_grad=0, device=cuda:0),
      %conv.bias : Float(64, strides=[1], requires_grad=0, device=cuda:0)):
  %/conv/Conv_output_0 : Float(1, 64, 16, 16, strides=[16384, 256, 16, 1], requires_grad=0, device=cuda:0) = onnx::Conv[dilations=[1, 1], group=1, kernel_shape=[3, 3], pads=[1, 1, 1, 1], strides=[1, 1], onnx_name="/conv/Conv"](%onnx::Conv_0, %conv.weight, %conv.bias), scope: __main__.MyModule::/torch.nn.modules.conv.Conv2d::conv # /home/litao/workspace/dl-frameworks-howto/venv/lib/python3.10/site-packages/torch/nn/modules/conv.py:458:0
  %/ReduceSum_output_0 : Float(requires_grad=1, device=cuda:0) = onnx::ReduceSum[keepdims=0, onnx_name="/ReduceSum"](%/conv/Conv_output_0), scope: __main__.MyModule:: # /home/litao/workspace/dl-frameworks-howto/./export_backward.py:9:0
  %/Constant_output_0 : Float(1, 4, 16, 16, strides=[1024, 256, 16, 1], requires_grad=0, device=cuda:0) = onnx::Constant[value=<Tensor>, onnx_name="/Constant"](), scope: __main__.MyModule:: # /home/litao/workspace/dl-frameworks-howto/./export_backward.py:10:0
  %6 : Float(1, 4, 16, 16, strides=[1024, 256, 16, 1], requires_grad=1, device=cuda:0) = onnx::Sub[onnx_name="/Sub"](%/ReduceSum_output_0, %/Constant_output_0), scope: __main__.MyModule:: # /home/litao/workspace/dl-frameworks-howto/./export_backward.py:10:0
  return (%6)

1.13.0+cu117
Traceback (most recent call last):
  File "/home/litao/workspace/dl-frameworks-howto/./export_backward.py", line 30, in <module>
    export_model('cpu')
  File "/home/litao/workspace/dl-frameworks-howto/./export_backward.py", line 23, in export_model
    torch.onnx.export(model, torch.ones([1,4,16,16]).to(device), f"conv-{device}.onnx", verbose=True)
  File "/home/litao/workspace/dl-frameworks-howto/venv/lib/python3.10/site-packages/torch/onnx/utils.py", line 504, in export
    _export(
  File "/home/litao/workspace/dl-frameworks-howto/venv/lib/python3.10/site-packages/torch/onnx/utils.py", line 1529, in _export
    graph, params_dict, torch_out = _model_to_graph(
  File "/home/litao/workspace/dl-frameworks-howto/venv/lib/python3.10/site-packages/torch/onnx/utils.py", line 1111, in _model_to_graph
    graph, params, torch_out, module = _create_jit_graph(model, args)
  File "/home/litao/workspace/dl-frameworks-howto/venv/lib/python3.10/site-packages/torch/onnx/utils.py", line 987, in _create_jit_graph
    graph, torch_out = _trace_and_get_graph_from_model(model, args)
  File "/home/litao/workspace/dl-frameworks-howto/venv/lib/python3.10/site-packages/torch/onnx/utils.py", line 891, in _trace_and_get_graph_from_model
    trace_graph, torch_out, inputs_states = torch.jit._get_trace_graph(
  File "/home/litao/workspace/dl-frameworks-howto/venv/lib/python3.10/site-packages/torch/jit/_trace.py", line 1184, in _get_trace_graph
    outs = ONNXTracedModule(f, strict, _force_outplace, return_inputs, _return_inputs_states)(*args, **kwargs)
  File "/home/litao/workspace/dl-frameworks-howto/venv/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1190, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/litao/workspace/dl-frameworks-howto/venv/lib/python3.10/site-packages/torch/jit/_trace.py", line 127, in forward
    graph, out = torch._C._create_graph_by_tracing(
  File "/home/litao/workspace/dl-frameworks-howto/venv/lib/python3.10/site-packages/torch/jit/_trace.py", line 118, in wrapper
    outs.append(self.inner(*trace_inputs))
  File "/home/litao/workspace/dl-frameworks-howto/venv/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1190, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/litao/workspace/dl-frameworks-howto/venv/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1178, in _slow_forward
    result = self.forward(*input, **kwargs)
  File "/home/litao/workspace/dl-frameworks-howto/./export_backward.py", line 11, in forward
    g = torch.autograd.grad(y, x)[0]
  File "/home/litao/workspace/dl-frameworks-howto/venv/lib/python3.10/site-packages/torch/autograd/__init__.py", line 300, in grad
    return Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
RuntimeError: Found an unsupported argument type in the JIT tracer. File a bug report.

```

### Versions

Collecting environment information...
PyTorch version: 1.13.0+cu117
Is debug build: False
CUDA used to build PyTorch: 11.7
ROCM used to build PyTorch: N/A

OS: Ubuntu 22.04.1 LTS (x86_64)
GCC version: (Ubuntu 11.3.0-1ubuntu1~22.04) 11.3.0
Clang version: 14.0.0-1ubuntu1
CMake version: version 3.22.1
Libc version: glibc-2.35

Python version: 3.10.6 (main, Nov  2 2022, 18:53:38) [GCC 11.3.0] (64-bit runtime)
Python platform: Linux-5.15.0-53-generic-x86_64-with-glibc2.35
Is CUDA available: True
CUDA runtime version: Could not collect
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: 
GPU 0: NVIDIA GeForce GTX 1060 6GB
GPU 1: NVIDIA GeForce RTX 2060 SUPER

Nvidia driver version: 520.61.05
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] mypy==0.991
[pip3] mypy-extensions==0.4.3
[pip3] numpy==1.23.4
[pip3] torch==1.13.0
[conda] Could not collect