Leak RAM when training 8 core TPU

Bug description

I am training on 8 TPU cores but the memory CPU Leak when the epoch ends.

Model

mobilenetV2

Data

cifar100

Env

pytorch_lightning==1.9.3 pytorch==1.13.0 code run on google colab

How to reproduce the bug

No response

Error messages and logs

# Error messages and logs here please

Environment

#- Lightning Component (e.g. Trainer, LightningModule, LightningApp, LightningWork, LightningFlow):
#- PyTorch Lightning Version (e.g., 1.5.0):
#- Lightning App Version (e.g., 0.5.2):
#- PyTorch Version (e.g., 2.0):
#- Python version (e.g., 3.9):
#- OS (e.g., Linux):
#- CUDA/cuDNN version:
#- GPU models and configuration:
#- How you installed Lightning(`conda`, `pip`, source):
#- Running environment of LightningApp (e.g. local, cloud):

More info

No response

cc @carmocca @JackCaoG @Liyang90 @gkroiz @borda @steventk-g