Repository: pytorch
Issue: Tracing doesn't work with spectral normalization
Labels: oncall: jit,weeks

## 🐛 Bug

When a module involves spectral normalization (using `nn.utils.spectral_norm`), traced module crashes on the first call. As a specific example, the attached script crashes.

## To Reproduce

Steps to reproduce the behavior:

1. Run the following code.

```python
import torch as th
from torch import nn, jit

mod = nn.Sequential(nn.utils.spectral_norm(nn.Linear(10, 100)), nn.Linear(100, 1))
inp = th.randn(1, 10)
print(mod(inp))
print(mod(inp))
# suppress check_trace warnings since SN is non-deterministic and mismatch is expected
traced = jit.trace_module(mod, {'forward': inp}, check_trace=False)
print(traced.forward(inp))
```

2. The output will be similar to

```
tensor([[0.1570]], grad_fn=<AddmmBackward>)
tensor([[0.1438]], grad_fn=<AddmmBackward>)
Traceback (most recent call last):
  File "a.py", line 10, in <module>
    print(traced.forward(inp))
RuntimeError: The following operation failed in the TorchScript interpreter.
Traceback of TorchScript (most recent call last):
/home/ziyu/anaconda3/envs/pyt17/lib/python3.6/site-packages/torch/nn/functional.py(3792): normalize
/home/ziyu/anaconda3/envs/pyt17/lib/python3.6/site-packages/torch/nn/utils/spectral_norm.py(84): compute_weight
/home/ziyu/anaconda3/envs/pyt17/lib/python3.6/site-packages/torch/nn/utils/spectral_norm.py(105): __call__
/home/ziyu/anaconda3/envs/pyt17/lib/python3.6/site-packages/torch/nn/modules/module.py(719): _call_impl
/home/ziyu/anaconda3/envs/pyt17/lib/python3.6/site-packages/torch/nn/modules/container.py(117): forward
/home/ziyu/anaconda3/envs/pyt17/lib/python3.6/site-packages/torch/nn/modules/module.py(709): _slow_forward
/home/ziyu/anaconda3/envs/pyt17/lib/python3.6/site-packages/torch/nn/modules/module.py(725): _call_impl
/home/ziyu/anaconda3/envs/pyt17/lib/python3.6/site-packages/torch/jit/_trace.py(940): trace_module
a.py(9): <module>
RuntimeError: div(): functions with out=... arguments don't support automatic differentiation, but one of the arguments requires grad.
```

## Expected behavior

The last `print` shouldn't crash, and should give similar result to the previous two.

## Environment

Please copy and paste the output from our
[environment collection script](https://raw.githubusercontent.com/pytorch/pytorch/master/torch/utils/collect_env.py)
(or fill out the checklist below manually).

```
Collecting environment information...
PyTorch version: 1.7.1
Is debug build: False
CUDA used to build PyTorch: 10.2
ROCM used to build PyTorch: N/A

OS: Ubuntu 18.04.4 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

Python version: 3.6 (64-bit runtime)
Is CUDA available: True
CUDA runtime version: 10.2.89
GPU models and configuration:
GPU 0: GeForce GTX 1080 Ti
GPU 1: GeForce GTX 1080 Ti
GPU 2: GeForce GTX 1080 Ti
GPU 3: GeForce GTX 1080 Ti

Nvidia driver version: 450.36.06
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] numpy==1.19.0
[conda] numpy                     1.18.4                    <pip>
[conda] pytorch-pretrained-biggan 0.1.1                     <pip>
[conda] torch                     1.7.1                     <pip>
```

## Additional context

Based on the error message and the content of `torch/nn/utils/spectral_norm.py(84): compute_weight`, it seems that a `no_grad` was not respected when calling the traced function.

cc @gmagogsfm