Repository: pytorch
Issue: Memory leak while training model generated by torch.fx.symbolic_trace() in data parallel mode
Labels: triaged,module: fx

### 🐛 Describe the bug

**There is memory leak when I train a model in data parallel mode if this model is generated by  torch.fx.symbolic_trace().**
**There seems to be no memory leak if I turn off one of the follow options:**
- data parallel mode (**`torch.nn.DataParallel`**)
- This model is **`torch.fx.GraphModule`** created via **`torch.fx.Tracer`**

The python script that can reproduce this BUG is attached here. This program would leak over 1GB memory per 10 training epochs.

```py
import torch
import torch.nn as nn
import torch.fx as fx
import torchvision
import torchvision.transforms as transforms

batch_size = 64

class MyNet(nn.Module):
    def __init__(self):
        super(MyNet, self).__init__()

        self.backbone = torchvision.models.densenet121()
        self.linear = nn.Linear(1000, 10)

    def forward(self, x):
        out = self.backbone(x)
        out = self.linear(out)

        return out

train_dataset = torchvision.datasets.CIFAR10(root='data/',
                                                train=True,
                                                transform=transforms.ToTensor(),
                                                download=True)

train_loader = torch.utils.data.DataLoader(dataset=train_dataset,
                                            batch_size=batch_size, 
                                            shuffle=True)

model = MyNet()
model = fx.symbolic_trace(model)

model = model.to('cuda')
model = nn.DataParallel(model)
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(
                model.parameters(),
                lr=0.01,
                momentum=0.8,
                nesterov=True,
                weight_decay=0.0001)

for epoch in range(1000):
    print('epoch:', epoch)
    for i, (images, labels) in enumerate(train_loader):

        images = images.to('cuda')
        labels = labels.to('cuda')
        
        outputs = model(images)

        loss = criterion(outputs, labels)
        
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        _, predicted = torch.max(outputs.reshape(-1, 10).data, 1)
        correctness = (predicted == labels).sum().item() / batch_size

        print('\ttraining loss: {:3.4f}'.format(loss.item()))
        print('\ttraining correctness: {:3.4f}'.format(correctness))
```

### Versions

```
PyTorch version: 1.10.0+cu102
Is debug build: False
CUDA used to build PyTorch: 10.2
ROCM used to build PyTorch: N/A

OS: Ubuntu 16.04.7 LTS (x86_64)
GCC version: (Ubuntu 5.4.0-6ubuntu1~16.04.12) 5.4.0 20160609
Clang version: Could not collect
CMake version: version 3.5.1
Libc version: glibc-2.17

Python version: 3.7.11 (default, Jul 27 2021, 14:32:16)  [GCC 7.5.0] (64-bit runtime)
Python platform: Linux-3.10.0-957.27.2.el7.x86_64-x86_64-with-debian-stretch-sid
Is CUDA available: True
CUDA runtime version: 10.2.89
GPU models and configuration: 
GPU 0: Tesla V100-SXM2-32GB
GPU 1: Tesla V100-SXM2-32GB

Nvidia driver version: 450.80.02
cuDNN version: /usr/lib/x86_64-linux-gnu/libcudnn.so.7.6.5
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] numpy==1.21.4
[pip3] torch==1.10.0
[pip3] torchaudio==0.10.0
[pip3] torchvision==0.11.1
[conda] blas                      1.0                         mkl  
[conda] mkl                       2021.4.0           h06a4308_640  
[conda] mkl-service               2.4.0            py37h7f8727e_0  
[conda] mkl_fft                   1.3.1            py37hd3c417c_0  
[conda] mkl_random                1.2.2            py37h51133e4_0  
[conda] numpy                     1.21.4                   pypi_0    pypi
[conda] numpy-base                1.21.2           py37h79a1101_0  
[conda] torch                     1.10.0                   pypi_0    pypi
[conda] torchaudio                0.10.0                   pypi_0    pypi
[conda] torchvision               0.11.1                   pypi_0    pypi
```

cc @ezyang @SherlockNoMad