Repository: pytorch
Issue: Temporarily disabling grad doesn't work if AMP is enabled
Labels: high priority,triaged,actionable,module: amp (automated mixed precision)

## 🐛 Bug

grad fails to engage properly after being temporarily disabled under AMP.

## To Reproduce

Here's some minimal reproduction code:

```
import torch
import torch.nn as nn

class RecycledLinear(nn.Module):
    def __init__(self, no_cycles):
        super(RecycledLinear, self).__init__()

        self.no_cycles = no_cycles
        self.linear = nn.Linear(10, 10)

    def forward(self, inp):
        out = inp
        # Run the model no_cycles times, enabling grad only on the last iteration
        for i in range(self.no_cycles):
            with torch.set_grad_enabled(i == (self.no_cycles - 1)):
                out = self.linear(out)

        return out


no_cycles = 2
x = torch.rand(10, 10, requires_grad=False).cuda()
model = RecycledLinear(no_cycles).cuda()

with torch.cuda.amp.autocast():
    out = model(x)

print(out.requires_grad) # prints False
```

The output emerges with grad disabled. This happens in spite of the fact that in the final iteration, `torch.grad.is_enabled()` and `self.linear.weight.requires_grad` are both true.

## Expected behavior

I would expect grad to be enabled on the output and for the operations in the final iteration of the network to have been recorded properly. This is what happens when the code is run without AMP.

## Environment

PyTorch version: 1.9.1+cu111
Is debug build: False
CUDA used to build PyTorch: 11.1
ROCM used to build PyTorch: N/A

OS: Ubuntu 18.04.6 LTS (x86_64)
GCC version: (Ubuntu 7.5.0-3ubuntu1~18.04) 7.5.0
Clang version: Could not collect
CMake version: version 3.10.2
Libc version: glibc-2.27

Python version: 3.9.5 (default, Jun  4 2021, 12:28:51)  [GCC 7.5.0] (64-bit runtime)
Python platform: Linux-4.15.0-58-generic-x86_64-with-glibc2.27
Is CUDA available: True
CUDA runtime version: 11.1.74
GPU models and configuration:
GPU 0: GeForce RTX 2080 Ti
GPU 1: GeForce RTX 2080 Ti
GPU 2: GeForce RTX 2080 Ti
GPU 3: GeForce RTX 2080 Ti
GPU 4: GeForce RTX 2080 Ti
GPU 5: GeForce RTX 2080 Ti
GPU 6: GeForce RTX 2080 Ti
GPU 7: GeForce RTX 2080 Ti

Nvidia driver version: 455.23.05
cuDNN version: /usr/lib/x86_64-linux-gnu/libcudnn.so.7.6.2
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] light-the-torch==0.3.5
[pip3] numpy==1.19.3
[pip3] pytorch-lightning==1.4.8
[pip3] torch==1.9.1+cu111
[pip3] torchmetrics==0.5.1
[conda] cudatoolkit               11.2.2               he111cf0_8    conda-forge
[conda] light-the-torch           0.3.5                    pypi_0    pypi
[conda] numpy                     1.19.3                   pypi_0    pypi
[conda] pytorch-lightning         1.4.8                    pypi_0    pypi
[conda] torch                     1.9.1+cu111              pypi_0    pypi
[conda] torchmetrics              0.5.1                    pypi_0    pypi

## Additional comments

Any tips for how to work around this issue would be greatly appreciated.

cc @ezyang @gchanan @zou3519 @kadeng @msaroufim @mcarilli @ptrblck @leslie-fang-intel @jgong5 @albanD @gqchen @pearu @nikitaved @soulitzer @Lezcano @Varal7