CPU memory consumption increases at the end of every epochs on TPU

Nobody has claimed this yet.

## 🐛 Bug
CPU memory consumption is stable during training.
But memory consumption keep increasing and never decreases at the end of every epochs.
It occurs when I use TPU. (It doesn't occur when I use GPU)

### To Reproduce

```python
import random
import numpy as np
import torch
import torch.nn as nn
import pytorch_lightning as pl
from torch.nn import TransformerEncoder, TransformerEncoderLayer
from torch.utils.data import Dataset, DataLoader

vocab_size = 5000

class TransformerModel(pl.LightningModule):
def __init__(self, ntoken=vocab_size, ninp=1024, nhead=2, nhid=1024, nlayers=2):
super(TransformerModel, self).__init__()
encoder_layers = TransformerEncoderLayer(ninp, nhead, nhid)
self.transformer_encoder = TransformerEncoder(encoder_layers, nlayers)
self.encoder = nn.Embedding(ntoken, ninp)
self.decoder = nn.Linear(ninp, ntoken)
self.criterion = nn.CrossEntropyLoss()

def forward(self, src):
src = self.encoder(src)
output = self.transformer_encoder(src)
output = self.decoder(output)
return output

def configure_optimizers(self):
optimizer = torch.optim.SGD(self.parameters(), lr=5.)
return optimizer

def training_step(self, batch, batch_idx, *args, **kwargs):
output = self(batch[0])
loss = self.criterion(output.transpose(1, 2), batch[0])
self.log('train_loss', loss)
return loss

class DummyDataset(Dataset):
def __init__(self, batch_size=256, bptt=32, size=1024):
self.data = [torch.randint(0, vocab_size, (batch_size, bptt)) for _ in range(size)]

def __len__(self):
return len(self.data)

def __getitem__(self, idx):
return self.data[idx]

random.seed(0)
np.random.seed(0)
torch.manual_seed(0)

dataset = DummyDataset()
data_loader = DataLoader(dataset)

model = TransformerModel(vocab_size)
trainer = pl.Trainer(gradient_clip_val=0.5, max_epochs=3, tpu_cores=8)
trainer.fit(model, data_loader)
```

### Expected behavior
Memory consumption should not be increased

### Environment
- PyTorch Lightning Version: 1.3.8
- torch-xla Version: 1.8
- PyTorch Version: 1.8.0
- Python version: 3.7.10
- OS: Debian GNU/Linux 10
- TPU type: v3-8
- TPU software version: pytorch-1.8
- TPU architecture: TPU node

### Additional context

cc @tchaton @rohitgr7 @akihironitta @kaushikb11

Start by running the provided TransformerModel and DataLoader reproduction through trainer.fit with the listed PyTorch Lightning, torch-xla, and TPU versions. Trace memory usage at epoch boundaries in the TPU training path; done means CPU memory remains stable across epochs without regressing GPU behavior.

Written by the indexing model from the issue text.