Repository: pytorch
Issue: Problems differentiating through certain operations when exporting to TorchScript
Labels: oncall: jit,triaged

### 🐛 Describe the bug

I'm trying to trace a model to TorchScript involving back propagation, and I've run into a number of issues. I can work around some of them, but this one seems a bit tough. You can see the kind of problem I've encountered if you run this python code:
```python
import torch

class Model(torch.nn.Module):

    def forward(self, inputs):
        xseq = inputs[0].detach().requires_grad_()
        with torch.enable_grad():
            # differentiating through this "loss" fails to export:
            loss = torch.matmul(xseq, xseq.transpose(-2, -1)).sum()
            # differentiating through this "loss" also fails to export:
            # loss = torch.sqrt(xseq).sum()
            # differentiating through this "loss" is fine:
            # loss = torch.pow(xseq, 0.5).sum()
            return torch.autograd.grad([loss], [xseq])[0]


torch.jit.trace(
    Model(),
    (torch.randn([2, 2, 10]) ** 2,),
)
```
As you can see, I've got a number of loss functions there that I'm computing gradients for, the final two of which are commented out. If I trace the model with the first or the second one uncommented, I get the following error:
```
RuntimeError: Cannot insert a Tensor that requires grad as a constant. Consider making it a parameter or input, or detaching the gradient
```
This clearly shouldn't happen as there basically aren't any constants involved in the computation (also happens if I make xseq square in the last 2 dimensions and remove the transpose(-2,-1)...).

Strangely, the export goes through fine if I uncomment the third loss function, although it should be equivalent to the second one.


### Versions

Collecting environment information...
PyTorch version: 2.2.1+cu118
Is debug build: False
CUDA used to build PyTorch: 11.8
ROCM used to build PyTorch: N/A

OS: Microsoft Windows 10 Enterprise
GCC version: Could not collect
Clang version: Could not collect
CMake version: version 3.25.1
Libc version: N/A

Python version: 3.10.11 (tags/v3.10.11:7d4cc5a, Apr  5 2023, 00:38:17) [MSC v.1929 64 bit (AMD64)] (64-bit runtime)
Python platform: Windows-10-10.0.19045-SP0
Is CUDA available: True
CUDA runtime version: 10.0.130
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: GPU 0: NVIDIA GeForce RTX 3090 Ti
Nvidia driver version: 536.23
cuDNN version: C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v11.8\bin\cudnn_ops_train64_8.dll
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture=9
CurrentClockSpeed=3501
DeviceID=CPU0
Family=107
L2CacheSize=16384
L2CacheSpeed=
Manufacturer=AuthenticAMD
MaxClockSpeed=3501
Name=AMD Ryzen Threadripper PRO 3975WX 32-Cores
ProcessorType=3
Revision=12544

Versions of relevant libraries:
[pip3] functorch==2.0.0
[pip3] lovely-numpy==0.2.8
[pip3] numpy==1.24.3
[pip3] onnx==1.14.1
[pip3] onnx-graphsurgeon==0.3.27
[pip3] onnxconverter-common==1.13.0
[pip3] onnxruntime==1.15.1
[pip3] optree==0.10.0
[pip3] pytorch-lightning==1.4.2
[pip3] tf2onnx==1.16.0
[pip3] torch==2.2.1+cu118
[pip3] torch-cluster==1.6.1
[pip3] torch-fidelity==0.3.0
[pip3] torch-geometric==2.3.0
[pip3] torch-scatter==2.1.1
[pip3] torch-sparse==0.6.17
[pip3] torch-spline-conv==1.2.2
[pip3] torchaudio==2.2.1+cu118
[pip3] torchmetrics==0.6.0
[pip3] torchvision==0.17.1+cu118
[conda] blas                      1.0                         mkl
[conda] cudatoolkit               11.6.0              hc0ea762_10    conda-forge
[conda] libblas                   3.9.0              16_win64_mkl    conda-forge
[conda] libcblas                  3.9.0              16_win64_mkl    conda-forge
[conda] liblapack                 3.9.0              16_win64_mkl    conda-forge
[conda] mkl                       2022.1.0           h6a75c08_874    conda-forge
[conda] mkl-include               2023.2.0            intel_49496    intel
[conda] mkl-static                2023.2.0            intel_49496    intel
[conda] numpy                     1.24.2          py310hd02465a_0    conda-forge
[conda] pytorch                   1.12.0          py3.10_cuda11.6_cudnn8_0    pytorch
[conda] pytorch-mutex             1.0                        cuda    pytorch

cc @EikanWang @jgong5 @wenzhe-nrv @sanchitintel