`No inf checks were recorded for this optimizer` when using SWA together with batch norm layers

Nobody has claimed this yet.

### Bug description

When SWA is used together with a model which has batch norm layers, the assertion `No inf checks were recorded for this optimizer.` is raised in the last epoch (=SWA epoch).

This worked fine with torch<2.0 but I am not sure whether it is a torch or lightning issue.

### How to reproduce the bug

```python
import os

import torch
from torch.utils.data import DataLoader, Dataset

from lightning import LightningModule, Trainer
from lightning.pytorch.callbacks import StochasticWeightAveraging
import torch.nn as nn

class RandomDataset(Dataset):
def __init__(self, size, length):
self.len = length
self.data = torch.randn(length, size)

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

def __len__(self):
return self.len

class BoringModel(LightningModule):
def __init__(self):
super().__init__()
self.layer = torch.nn.Linear(100, 20)

# This line is the problem even though the batch norm layer is not even used
self.norm = nn.BatchNorm1d(10)

self.ce_loss = nn.CrossEntropyLoss(weight=torch.ones(20))

def forward(self, x):
return self.layer(x)

def training_step(self, batch, batch_idx):
loss = self(batch).sum()
self.log("train_loss", loss)
return {"loss": loss}

def validation_step(self, batch, batch_idx):
loss = self.ce_loss(self(batch), self(batch))
self.log("valid_loss", loss)

def test_step(self, batch, batch_idx):
loss = self(batch).sum()
self.log("test_loss", loss)

def configure_optimizers(self):
return torch.optim.Adam(self.layer.parameters(), lr=0.1)

def run():
train_data = DataLoader(RandomDataset(100, 64), batch_size=2)
val_data = DataLoader(RandomDataset(100, 64), batch_size=2)

model = BoringModel()
trainer = Trainer(
default_root_dir=os.getcwd(),
limit_train_batches=2,
limit_val_batches=1,
limit_test_batches=1,
num_sanity_val_steps=0,
max_epochs=10,
enable_model_summary=False,
callbacks=[StochasticWeightAveraging(swa_lrs=1e-2)],
precision=16,
accelerator='gpu',
devices=1,
)
trainer.fit(model, train_dataloaders=train_data, val_dataloaders=val_data)

if __name__ == "__main__":
run()
```

### Error messages and logs

```
Traceback (most recent call last):
File "/mnt/ssd_8tb/htc/src/tt.py", line 75, in
run()
File "/mnt/ssd_8tb/htc/src/tt.py", line 70, in run
trainer.fit(model, train_dataloaders=train_data, val_dataloaders=val_data)
File "/home/j562r/miniconda3/envs/htc2/lib/python3.10/site-packages/lightning/pytorch/trainer/trainer.py", line 520, in fit
call._call_and_handle_interrupt(
File "/home/j562r/miniconda3/envs/htc2/lib/python3.10/site-packages/lightning/pytorch/trainer/call.py", line 44, in _call_and_handle_interrupt
return trainer_fn(*args, **kwargs)
File "/home/j562r/miniconda3/envs/htc2/lib/python3.10/site-packages/lightning/pytorch/trainer/trainer.py", line 559, in _fit_impl
self._run(model, ckpt_path=ckpt_path)
File "/home/j562r/miniconda3/envs/htc2/lib/python3.10/site-packages/lightning/pytorch/trainer/trainer.py", line 935, in _run
results = self._run_stage()
File "/home/j562r/miniconda3/envs/htc2/lib/python3.10/site-packages/lightning/pytorch/trainer/trainer.py", line 978, in _run_stage
self.fit_loop.run()
File "/home/j562r/miniconda3/envs/htc2/lib/python3.10/site-packages/lightning/pytorch/loops/fit_loop.py", line 201, in run
self.advance()
File "/home/j562r/miniconda3/envs/htc2/lib/python3.10/site-packages/lightning/pytorch/loops/fit_loop.py", line 354, in advance
self.epoch_loop.run(self._data_fetcher)
File "/home/j562r/miniconda3/envs/htc2/lib/python3.10/site-packages/lightning/pytorch/loops/training_epoch_loop.py", line 133, in run
self.advance(data_fetcher)
File "/home/j562r/miniconda3/envs/htc2/lib/python3.10/site-packages/lightning/pytorch/loops/training_epoch_loop.py", line 218, in advance
batch_output = self.automatic_optimization.run(trainer.optimizers[0], kwargs)
File "/home/j562r/miniconda3/envs/htc2/lib/python3.10/site-packages/lightning/pytorch/loops/optimization/automatic.py", line 185, in run
self._optimizer_step(kwargs.get("batch_idx", 0), closure)
File "/home/j562r/miniconda3/envs/htc2/lib/python3.10/site-packages/lightning/pytorch/loops/optimization/automatic.py", line 261, in _optimizer_step
call._call_lightning_module_hook(
File "/home/j562r/miniconda3/envs/htc2/lib/python3.10/site-packages/lightning/pytorch/trainer/call.py", line 142, in _call_lightning_module_hook
output = fn(*args, **kwargs)
File "/home/j562r/miniconda3/envs/htc2/lib/python3.10/site-packages/lightning/pytorch/core/module.py", line 1265, in optimizer_step
optimizer.step(closure=optimizer_closure)
File "/home/j562r/miniconda3/envs/htc2/lib/python3.10/site-packages/lightning/pytorch/core/optimizer.py", line 158, in step
step_output = self._strategy.optimizer_step(self._optimizer, closure, **kwargs)
File "/home/j562r/miniconda3/envs/htc2/lib/python3.10/site-packages/lightning/pytorch/strategies/strategy.py", line 224, in optimizer_step
return self.precision_plugin.optimizer_step(optimizer, model=model, closure=closure, **kwargs)
File "/home/j562r/miniconda3/envs/htc2/lib/python3.10/site-packages/lightning/pytorch/plugins/precision/amp.py", line 83, in optimizer_step
step_output = self.scaler.step(optimizer, **kwargs)
File "/home/j562r/miniconda3/envs/htc2/lib/python3.10/site-packages/torch/cuda/amp/grad_scaler.py", line 370, in step
assert len(optimizer_state["found_inf_per_device"]) > 0, "No inf checks were recorded for this optimizer."
AssertionError: No inf checks were recorded for this optimizer.
Epoch 10: 67%|██████▋ | 2/3 [00:00<00:00, 14.17it/s, v_num=0]
```

### Environment

Current environment

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

### More info

_No response_

Reproduce the failure with the supplied SWA, batch-normalization, and mixed-precision example, then inspect the optimizer-step path in lightning/pytorch/plugins/precision/amp.py and the surrounding optimizer code named in the traceback. Add regression coverage for the final SWA epoch and verify training completes without the assertion.

Written by the indexing model from the issue text.