Repository: pytorch
Issue: Pruning `torch.nn.MultiheadAttention` causes RuntimeError
Labels: triaged,module: pruning,Stale

## 🐛 Bug

I am running into the following RuntimeError when pruning parameters of `torch.nn.MultiheadAttention` module:
```
RuntimeError: Trying to backward through the graph a second time (or directly access saved tensors after they have already been freed). Saved intermediate values of the graph are freed when you call .backward() or autograd.grad(). Specify retain_graph=True if you need to backward through the graph a second time or if you need to access saved tensors after calling backward.
```

## To Reproduce

Steps to reproduce the behavior:

Here is the toy example I'm running:
```
import torch.nn as nn
import torch
import torch.nn.utils.prune as prune

class net(nn.Module):
    def __init__(self):
        super().__init__()

        self.att = nn.MultiheadAttention(
            embed_dim=100,
            num_heads=2,
            batch_first=True
        )
        self.proj = nn.Linear(100, 10)
        self.criterion = nn.CrossEntropyLoss()

    def forward(self, x):
        attn_output, attn_weights = self.att(x, x, x)
        logits = self.proj(attn_output)
        return logits

    def compute_loss(self, logits, target):
        loss = self.criterion(logits, target)
        return loss

model = net().cuda()
# prune here
params_to_prune = [
    (model.att, "in_proj_weight"),
    (model.att.out_proj, "weight"),
    (model.proj, "weight")
]
prune.global_unstructured(
    params_to_prune,
    pruning_method=prune.L1Unstructured,
    amount=0.05,
)

opt = torch.optim.SGD(model.parameters(), lr=0.01)
for step in range(0, 10):
    x = torch.rand((16, 50, 100)).cuda()  # data of batchsize=16, seqlength=50, featuredim=100
    y = torch.ones((16*50, )).long().cuda()
    opt.zero_grad()
    logits = model(x)
    loss = model.compute_loss(logits.reshape(16*50, 10), y)
    loss.backward()
    opt.step()
    print("step {} successful".format(step))

```

Full error stack:
```
step 0 successful
Traceback (most recent call last):
  File "test_MHA_prune.py", line 49, in <module>
    loss.backward()
  File "/scratch2/mxy190019/anaconda3/envs/custom-pytorch/lib/python3.8/site-packages/torch/_tensor.py", line 307, in backward
    torch.autograd.backward(self, gradient, retain_graph, create_graph, inputs=inputs)
  File "/scratch2/mxy190019/anaconda3/envs/custom-pytorch/lib/python3.8/site-packages/torch/autograd/__init__.py", line 154, in backward
    Variable._execution_engine.run_backward(
RuntimeError: Trying to backward through the graph a second time (or directly access saved tensors after they have already been freed). Saved intermediate values of the graph are freed when you call .backward() or autograd.grad(). Specify retain_graph=True if you need to backward through the graph a second time or if you need to access saved tensors after calling backward.
```

## Expected behavior
The pruned model should be able to train without a problem.

## Environment

 - PyTorch Version (e.g., 1.0): 1.10.0+cu102
 - OS (e.g., Linux): CentOS Linux 7 (Core) (x86_64)
 - How you installed PyTorch (`conda`, `pip`, source): `pip install torch`
 - Build command you used (if compiling from source): N/A
 - Python version: 3.8.0 (default, Nov  6 2019, 21:49:08)  [GCC 7.3.0] (64-bit runtime)
 - CUDA/cuDNN version: CUDA 10.2
 - GPU models and configuration: One single GeForce RTX 2080 Ti
 - Any other relevant information: None

## Additional context

It seems that the error is brought by the pruning of this parameter in `torch.nn.MultiheadAttention`: (model.att.out_proj, "weight"). If I remove this line then there is no such RuntimeError. Am I missing anything here for pruning the MultiheadAttention module? How should I properly prune the parameters in `out_proj` (which I believe is just a Linear layer)? 

Thanks!


cc @jbschlosser @bhosmer @cpuhrsch