Repository: pytorch
Issue: torch.cuda.empty_cache() doesn't work correctly for the first instance of nn.Module set to cuda (and a potential method to deal with that).
Labels: needs reproduction,module: cuda,module: memory usage,triaged

### 🐛 Describe the bug

I am using Google Colab with T4 runtime type. The issue is, torch.cuda.empty_cache() cannot clear the ram on GPU for the first instance of nn.Moudle set to cuda. This also happened when I used other gpu cloud services.

I did several simple experiments. Below are the code I will run each time I restart the session (Base Code):

```python
device = "cuda:0"
import gc
import torch
import torch.nn as nn
from torchvision.models import vit_b_16

class Decoder(nn.Module):
    def __init__(self, latent_dim, hidden_dim, output_dim):
        super(Decoder, self).__init__()

        self.decoder = nn.Sequential(
            nn.Linear(latent_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, output_dim)
        )

    def forward(self, x):
        return self.decoder(x)

class AutoEncoder(nn.Module):
    def __init__(self, encoder, decoder):
        super(AutoEncoder, self).__init__()
        self.encoder = encoder
        self.decoder = decoder

    def forward(self, x, mask):
        x_masked = mask * x
        encode_patch = self.encoder(x_masked)
        reconstruct = self.decoder(encode_patch)

        return reconstruct
```

# Exp 1: One instance
Restart the sesson, and run the Base Code. The, run the code below:
```python
encoder = vit_b_16(pretrained=True)
decoder = Decoder(1000, 512, 3*224*224)
ae_model = AutoEncoder(encoder, decoder).to(device)
```
Now, the GPU usage becomes *0.8* G.

Then run the code:
```python
# model = ae_model.to('cpu')
del ae_model
# gc.collect()
torch.cuda.empty_cache()
``` 
torch.cuda.empty_cache() does't work. The GPU usage is still *0.8* G.

# Exp 2: One instance + gc.collect()
In this Experiment, I try gc.collect().

As usual, restart the session and run the Base Code.

Then run the code below:
```python
encoder = vit_b_16(pretrained=True)
decoder = Decoder(1000, 512, 3*224*224)
ae_model = AutoEncoder(encoder, decoder).to(device)
```

```python
# model = ae_model.to('cpu')
del ae_model
gc.collect()
torch.cuda.empty_cache()
```
The GPU usage is still *0.8* G, i.e., gc.collect() does't help to clear the memory on GPU.

# Exp 3: Two seperate instances 

In this experiment, I use two seperate instances of nn.Moudle set to cuda.

Restart the session and run the Base Code.

Then run the code below:
```python
encoder = vit_b_16(pretrained=True)
decoder = Decoder(1000, 512, 3*224*224)
ae_model = AutoEncoder(encoder, decoder).to(device)
```

```python
# model = ae_model.to('cpu')
del ae_model
# gc.collect()
torch.cuda.empty_cache()
```
(Yes, as you think, the GPU usage is still *0.8* G)

Now, create a new sperate instance and set it to cuda:

```python
encoder = vit_b_16(pretrained=True)
decoder = Decoder(1000, 512, 3*224*224)
ae_model = AutoEncoder(encoder, decoder).to(device)
```
Now the gpu the GPU usage becomes *1.0* G.

Then run the code:
```python
# model = ae_model.to('cpu')
del ae_model
# gc.collect()
torch.cuda.empty_cache()
```
Differently, this time, the second instance is cleared and the GPU usage becomes *0.8* G.

# Exp 4: One instance + set it to cpu
In this experiment, I set the model to CPU before call torch.cuda.empty_cache().

Restart the session and run the Base Code.

Then run the code below:
```python
encoder = vit_b_16(pretrained=True)
decoder = Decoder(1000, 512, 3*224*224)
ae_model = AutoEncoder(encoder, decoder).to(device)
```

The GPU usage is *0.8*G.

Now, run the code below:
```python
model = ae_model.to('cpu')
del ae_model
# gc.collect()
torch.cuda.empty_cache()
```

Suppressingly, this time, the memory is successfully cleared, and the GPU usage becomes *0.1*G.

# Conclusion
## For the first instance of nn.Module set to 'cuda'
You cannot clear cache directly using torch.cuda.empty_cache() when the reference count of the instance becomes 0, even if you use gc.collect()

How to deal with that (or how to clear the memory)? Set it to CPU before calling torch.cuda.empty_cache().
## For the second seperate instance of  nn.Module set to 'cuda'
torch.cuda.empty_cache() does work.

### Versions

--2024-07-15 11:01:54--  https://raw.githubusercontent.com/pytorch/pytorch/main/torch/utils/collect_env.py
Resolving raw.githubusercontent.com (raw.githubusercontent.com)... 185.199.108.133, 185.199.109.133, 185.199.110.133, ...
Connecting to raw.githubusercontent.com (raw.githubusercontent.com)|185.199.108.133|:443... connected.
HTTP request sent, awaiting response... 200 OK
Length: 23357 (23K) [text/plain]
Saving to: ‘collect_env.py’

collect_env.py      100%[===================>]  22.81K  --.-KB/s    in 0s      

2024-07-15 11:01:54 (135 MB/s) - ‘collect_env.py’ saved [23357/23357]

Collecting environment information...
PyTorch version: 2.3.0+cu121
Is debug build: False
CUDA used to build PyTorch: 12.1
ROCM used to build PyTorch: N/A

OS: Ubuntu 22.04.3 LTS (x86_64)
GCC version: (Ubuntu 11.4.0-1ubuntu1~22.04) 11.4.0
Clang version: 14.0.0-1ubuntu1.1
CMake version: version 3.27.9
Libc version: glibc-2.35

Python version: 3.10.12 (main, Mar 22 2024, 16:50:05) [GCC 11.4.0] (64-bit runtime)
Python platform: Linux-6.1.85+-x86_64-with-glibc2.35
Is CUDA available: True
CUDA runtime version: 12.2.140
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: GPU 0: Tesla T4
Nvidia driver version: 535.104.05
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.8.9.6
/usr/lib/x86_64-linux-gnu/libcudnn_adv_infer.so.8.9.6
/usr/lib/x86_64-linux-gnu/libcudnn_adv_train.so.8.9.6
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_infer.so.8.9.6
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_train.so.8.9.6
/usr/lib/x86_64-linux-gnu/libcudnn_ops_infer.so.8.9.6
/usr/lib/x86_64-linux-gnu/libcudnn_ops_train.so.8.9.6
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture:                         x86_64
CPU op-mode(s):                       32-bit, 64-bit
Address sizes:                        46 bits physical, 48 bits virtual
Byte Order:                           Little Endian
CPU(s):                               2
On-line CPU(s) list:                  0,1
Vendor ID:                            GenuineIntel
Model name:                           Intel(R) Xeon(R) CPU @ 2.00GHz
CPU family:                           6
Model:                                85
Thread(s) per core:                   2
Core(s) per socket:                   1
Socket(s):                            1
Stepping:                             3
BogoMIPS:                             4000.28
Flags:                                fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ss ht syscall nx pdpe1gb rdtscp lm constant_tsc rep_good nopl xtopology nonstop_tsc cpuid tsc_known_freq pni pclmulqdq ssse3 fma cx16 pcid sse4_1 sse4_2 x2apic movbe popcnt aes xsave avx f16c rdrand hypervisor lahf_lm abm 3dnowprefetch invpcid_single ssbd ibrs ibpb stibp fsgsbase tsc_adjust bmi1 hle avx2 smep bmi2 erms invpcid rtm mpx avx512f avx512dq rdseed adx smap clflushopt clwb avx512cd avx512bw avx512vl xsaveopt xsavec xgetbv1 xsaves arat md_clear arch_capabilities
Hypervisor vendor:                    KVM
Virtualization type:                  full
L1d cache:                            32 KiB (1 instance)
L1i cache:                            32 KiB (1 instance)
L2 cache:                             1 MiB (1 instance)
L3 cache:                             38.5 MiB (1 instance)
NUMA node(s):                         1
NUMA node0 CPU(s):                    0,1
Vulnerability Gather data sampling:   Not affected
Vulnerability Itlb multihit:          Not affected
Vulnerability L1tf:                   Mitigation; PTE Inversion
Vulnerability Mds:                    Vulnerable; SMT Host state unknown
Vulnerability Meltdown:               Vulnerable
Vulnerability Mmio stale data:        Vulnerable
Vulnerability Reg file data sampling: Not affected
Vulnerability Retbleed:               Vulnerable
Vulnerability Spec rstack overflow:   Not affected
Vulnerability Spec store bypass:      Vulnerable
Vulnerability Spectre v1:             Vulnerable: __user pointer sanitization and usercopy barriers only; no swapgs barriers
Vulnerability Spectre v2:             Vulnerable; IBPB: disabled; STIBP: disabled; PBRSB-eIBRS: Not affected; BHI: Vulnerable (Syscall hardening enabled)
Vulnerability Srbds:                  Not affected
Vulnerability Tsx async abort:        Vulnerable

Versions of relevant libraries:
[pip3] numpy==1.25.2
[pip3] torch==2.3.0+cu121
[pip3] torchaudio==2.3.0+cu121
[pip3] torchsummary==1.5.1
[pip3] torchtext==0.18.0
[pip3] torchvision==0.18.0+cu121
[pip3] triton==2.3.0
[conda] Could not collect

cc @ptrblck @msaroufim