Repository: pytorch
Issue: torch.jit.trace memory leak
Labels: oncall: jit

## 🐛 Bug

torch.jit.trace has memory leak.

## To Reproduce

Run following code

        import gc
        import os
        import sys
        
        import torch
        import psutil
        
        nFeatures = 52
        
        
        layer_neurons = [nFeatures,128,256,512,256,128,128,64,64,32,32,16,16,8,8]
        layers = []
        for i, o in zip(layer_neurons[:-1], layer_neurons[1:]):
            layers.append(torch.nn.Linear(i, o))
            layers.append(torch.nn.ReLU())
        layers.append(torch.nn.Linear(layer_neurons[-1], 3))
        net = torch.nn.Sequential(*layers).float().cuda()
        
        optimizer = torch.optim.Adam(net.parameters(), lr=0.00025)
        loss_fn = torch.nn.MSELoss()
        
        
        if __name__ == '__main__':
            print("Python version:", sys.version)
            print("PyTorch version: ", torch.__version__)
        
            process = psutil.Process(os.getpid())
        
            tensor = torch.randn(1, nFeatures).cuda()
        
            for i in range(2000):
                m = torch.jit.trace(net, tensor)
                del m
                if (i % 200) == 0:
                    gc.collect()
                    print(i,  "Memory", process.memory_info()[0])

Output captured:

        Python version: 3.8.8 (default, Feb 24 2021, 21:46:12) 
        [GCC 7.3.0]
        PyTorch version:  1.8.1+cu111
        0 Memory 4353466368
        200 Memory 4736237568
        400 Memory 5147541504
        600 Memory 5557080064
        800 Memory 5966946304
        1000 Memory 6375677952
        1200 Memory 6785871872
        1400 Memory 7196590080
        1600 Memory 7607848960
        1800 Memory 8016650240

## Expected behavior

<!-- A clear and concise description of what you expected to happen. -->

## Environment

        Collecting environment information...
        PyTorch version: 1.8.1+cu111
        Is debug build: False
        CUDA used to build PyTorch: 11.1
        ROCM used to build PyTorch: N/A
        
        OS: Ubuntu 20.04.2 LTS (x86_64)
        GCC version: (Ubuntu 9.3.0-17ubuntu1~20.04) 9.3.0
        Clang version: Could not collect
        CMake version: version 3.16.3
        
        Python version: 3.8 (64-bit runtime)
        Is CUDA available: True
        CUDA runtime version: 11.3.58
        GPU models and configuration: GPU 0: GeForce RTX 2060 SUPER
        Nvidia driver version: 460.73.01
        cuDNN version: Probably one of the following:
        /usr/lib/x86_64-linux-gnu/libcudnn.so.8.1.1
        /usr/lib/x86_64-linux-gnu/libcudnn_adv_infer.so.8.1.1
        /usr/lib/x86_64-linux-gnu/libcudnn_adv_train.so.8.1.1
        /usr/lib/x86_64-linux-gnu/libcudnn_cnn_infer.so.8.1.1
        /usr/lib/x86_64-linux-gnu/libcudnn_cnn_train.so.8.1.1
        /usr/lib/x86_64-linux-gnu/libcudnn_ops_infer.so.8.1.1
        /usr/lib/x86_64-linux-gnu/libcudnn_ops_train.so.8.1.1
        /usr/local/cuda-11.3/targets/x86_64-linux/lib/libcudnn.so.8.2.0
        /usr/local/cuda-11.3/targets/x86_64-linux/lib/libcudnn_adv_infer.so.8.2.0
        /usr/local/cuda-11.3/targets/x86_64-linux/lib/libcudnn_adv_train.so.8.2.0
        /usr/local/cuda-11.3/targets/x86_64-linux/lib/libcudnn_cnn_infer.so.8.2.0
        /usr/local/cuda-11.3/targets/x86_64-linux/lib/libcudnn_cnn_train.so.8.2.0
        /usr/local/cuda-11.3/targets/x86_64-linux/lib/libcudnn_ops_infer.so.8.2.0
        /usr/local/cuda-11.3/targets/x86_64-linux/lib/libcudnn_ops_train.so.8.2.0
        HIP runtime version: N/A
        MIOpen runtime version: N/A
        
        Versions of relevant libraries:
        [pip3] numpy==1.19.5
        [pip3] numpydoc==1.1.0
        [pip3] torch==1.8.1+cu111
        [pip3] torchaudio==0.8.1
        [pip3] torchvision==0.9.1+cu111
        [conda] blas                      1.0                         mkl  
        [conda] mkl                       2020.2                      256  
        [conda] mkl-service               2.3.0            py38he904b0f_0  
        [conda] mkl_fft                   1.3.0            py38h54f3939_0  
        [conda] mkl_random                1.1.1            py38h0573a6f_0  
        [conda] numpy                     1.18.5                   pypi_0    pypi
        [conda] numpydoc                  1.1.0              pyhd3eb1b0_1  
        [conda] torch                     1.8.1+cu111              pypi_0    pypi
        [conda] torchaudio                0.8.1                    pypi_0    pypi
        [conda] torchvision               0.9.1+cu111              pypi_0    pypi



cc @gmagogsfm