Repository: pytorch
Issue: torch.set_grad_enabled results in RuntimeError with torch.jit.script
Labels: oncall: jit

### 🐛 Describe the bug

Calling `torch.set_grad_enabled` in `nn.Module.forward` results in RuntimeError when compiling a module with `torch.jit.script`. 

Minimal Example:
```python
import torch
import torch.nn as nn


class BrokenModule(nn.Module):
    def __init__(self, freeze: bool) -> None:
        super().__init__()
        self.freeze = freeze

        self.maybe_frozen = nn.Linear(10, 10)
        self.never_frozen = nn.Linear(10, 10)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        with torch.set_grad_enabled(not self.freeze):
            x = self.maybe_frozen(x)
        x = self.never_frozen(x)
        return x

# This raises RuntimeError
script = torch.jit.script(BrokenModule(freeze=True))  # type: ignore
```

The full error message:
```
Traceback (most recent call last):
  File "/home/xxx/repos/yyy/torchbug.py", line 39, in <module>
    script = torch.jit.script(BrokenModule(freeze=True))  # type: ignore
  File "/home/xxx/venvs/dev38/lib/python3.8/site-packages/torch/jit/_script.py", line 1286, in script
    return torch.jit._recursive.create_script_module(
  File "/home/xxx/venvs/dev38/lib/python3.8/site-packages/torch/jit/_recursive.py", line 458, in create_script_module
    return create_script_module_impl(nn_module, concrete_type, stubs_fn)
  File "/home/lxxx/venvs/dev38/lib/python3.8/site-packages/torch/jit/_recursive.py", line 524, in create_script_module_impl
    create_methods_and_properties_from_stubs(concrete_type, method_stubs, property_stubs)
  File "/home/xxx/venvs/dev38/lib/python3.8/site-packages/torch/jit/_recursive.py", line 375, in create_methods_and_properties_from_stubs
    concrete_type._create_methods_and_properties(property_defs, property_rcbs, method_defs, method_rcbs, method_defaults)
RuntimeError: 
With item expression must return an object:
  File "/home/xxx/repos/yyy/torchbug.py", line 14
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        with torch.set_grad_enabled(not self.freeze):
             ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
            x = self.maybe_frozen(x)
        x = self.never_frozen(x)
```

Ugly Workaround:
```python
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if self.freeze:
            with torch.no_grad():
                x = self.maybe_frozen(x)
        else:
            x = self.maybe_frozen(x)
        x = self.never_frozen(x)
        return x
```
Why can't I just use the original script?

### Versions

PyTorch version: 1.12.0+cu116
Is debug build: False
CUDA used to build PyTorch: 11.6
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: Could not collect
CMake version: version 3.24.3
Libc version: glibc-2.31

Python version: 3.8.10 (default, Jun 22 2022, 20:18:18)  [GCC 9.4.0] (64-bit runtime)
Python platform: Linux-5.15.0-52-generic-x86_64-with-glibc2.29
Is CUDA available: True
CUDA runtime version: 11.6.124
CUDA_MODULE_LOADING set to: 
GPU models and configuration: GPU 0: NVIDIA GeForce RTX 2080 Ti
Nvidia driver version: 510.85.02
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-extensions==0.4.3
[pip3] numpy==1.23.1
[pip3] numpy-quaternion==2022.4.2
[pip3] pytorch-lightning==1.6.5
[pip3] torch==1.12.0+cu116
[pip3] torchmetrics==0.9.3
[pip3] torchvision==0.13.0+cu116
[conda] Could not collect

cc @EikanWang @jgong5 @wenzhe-nrv @sanchitintel