Repository: pytorch
Issue: [JIT] .backward() not supported by JIT trace
Labels: oncall: jit

# 🐛 Describe the bug
I want to export my pytorch model to onnx format. MyModel wraps an existing model and, for a given input, computes the wrapped model’s output and its gradient to the input. (Virtual env: Python 3.9, CUDA 11.7.1 and, PyTorch 1.13.1.)

The following code achieves what I want within PyTorch:
cc. @t-vi, @apaszke, @fritzo.

Code example
The following example:
```
import torch
import torch.nn as nn


#  MyModel we want to wrap
class BaseModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(in_channels=2, out_channels=1, kernel_size=3, padding=1)
    
    def forward(self, x):
        return self.conv1(x)

# the MyModel we want to export using ONNX
class MyModel(nn.Module):
    def __init__(self, forward_model):
        super().__init__()
        self.forward_model = forward_model
    
    def forward(self, x, dy):
        y = self.forward_model(x)
        y.backward(gradient=dy)
        dx = x.grad
        return y, dx


model = BaseModel()
MyModel = MyModel(model)

x = torch.randn(1, 2, 5, 5)  # input
dy = torch. Ones(1, 1, 5, 5)  # output  

# compute the wrapped model's output and the input sensitivity
# (i.e., the gradient w.r.t. the input and the given output sensitivity)
y, dx = MyModel(x, y)

torch.onnx.export(MyModel, (x, dy), "MyModel.onnx", input_names=["x", "dy"], output_names=["y", "dx"])
```

throws an error on a recent PyTorch version:
```
Traceback (most recent call last):
  File "MyModel.py", line 117, in <module>
    torch.onnx.export(MyModel, (x, dy), "MyModel.onnx", input_names=["x", "dy"], output_names=["y", "dx"])
  File "~/my-venv/lib/python3.9/site-packages/torch/onnx/__init__.py", line 316, in export
    return utils.export(model, args, f, export_params, verbose, training,
  File "~/my-venv/lib/python3.9/site-packages/torch/onnx/utils.py", line 107, in export
    _export(model, args, f, export_params, verbose, training, input_names, output_names,
  File "~/my-venv/lib/python3.9/site-packages/torch/onnx/utils.py", line 724, in _export
    _model_to_graph(model, args, verbose, input_names,
  File "~/my-venv/lib/python3.9/site-packages/torch/onnx/utils.py", line 493, in _model_to_graph
    graph, params, torch_out, module = _create_jit_graph(model, args)
  File "~/my-venv/lib/python3.9/site-packages/torch/onnx/utils.py", line 437, in _create_jit_graph
    graph, torch_out = _trace_and_get_graph_from_model(model, args)
  File "~/my-venv/lib/python3.9/site-packages/torch/onnx/utils.py", line 388, in _trace_and_get_graph_from_model
    torch.jit._get_trace_graph(model, args, strict=False, _force_outplace=False, _return_inputs_states=True)
  File "~/my-venv/lib/python3.9/site-packages/torch/jit/_trace.py", line 1166, in _get_trace_graph
    outs = ONNXTracedModule(f, strict, _force_outplace, return_inputs, _return_inputs_states)(*args, **kwargs)
  File "~/my-venv/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1102, in _call_impl
    return forward_call(*input, **kwargs)
  File "~/my-venv/lib/python3.9/site-packages/torch/jit/_trace.py", line 127, in forward
    graph, out = torch._C._create_graph_by_tracing(
  File "~/my-venv/lib/python3.9/site-packages/torch/jit/_trace.py", line 118, in wrapper
    outs.append(self. Inner(*trace_inputs))
  File "~/my-venv/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1102, in _call_impl
    return forward_call(*input, **kwargs)
  File "~/my-venv/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1090, in _slow_forward
    result = self. Forward(*input, **kwargs)
  File "pullback_onnx.py", line 94, in forward
    y.backward(gradient=dy)
  File "~/my-venv/lib/python3.9/site-packages/torch/_tensor.py", line 307, in backward
    torch.autograd.backward(self, gradient, retain_graph, create_graph, inputs=inputs)
  File "~/my-venv/lib/python3.9/site-packages/torch/autograd/__init__.py", line 154, in backward
    Variable._execution_engine.run_backward(
RuntimeError: Found an unsupported argument type in the JIT tracer. File a bug report.
```


# Versions

```
Collecting environment information...
PyTorch version: 1.13.1+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: Could not collect
CMake version: Could not collect
Libc version: glibc-2.35

Python version: 3.9.16 (main, Dec  6 2022, 18:36:13)  [GCC 11.3.0] (64-bit runtime)
Python platform: Linux-5.15.0-56-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 1660 Ti
Nvidia driver version: 525.60.13
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.5
[pip3] pytorch-lightning==1.8.4.post0
[pip3] torch==1.13.1
[pip3] torch-model-archiver==0.6.1
[pip3] torch-ort==1.13.1
[pip3] torch-workflow-archiver==0.2.5
[pip3] torchaudio==0.13.1
[pip3] torchinfo==1.7.1
[pip3] torchmetrics==0.10.0
[pip3] torchserve==0.6.1
[pip3] torchtext==0.14.0
[pip3] torchvision==0.14.1
[conda] Could not collect

```

cc @EikanWang @jgong5 @wenzhe-nrv @sanchitintel