Callbacks leaking test gradients?

Nobody has claimed this yet.

### Bug description

Hi, I have the following module:

```python
class SpatialTransformer(ContinualModule):
"""A model that combines a parameter regressor and differentiable affine transforms."""

def __init__(
self,
img_size: int = 64,
in_channels: int = 1,
channels: Optional[list] = None,
gamma: float = 0.5,
**kwargs,
):
super().__init__(**kwargs)

if channels is None:
channels = [16, 16, 32, 32, 64]
self.encoder = Encoder(channels, in_channels)
self.encoder_output_dim = (img_size // 2 ** len(channels)) ** 2 * channels[-1]
self.regressor = MLP(
dims=[self.encoder_output_dim, 64, 32, 6],
)
# initialize the regressor to the identity transform
self.regressor.model[-1].weight.data.zero_()
self.regressor.model[-1].bias.data.copy_(
torch.tensor([1, 0, 0, 0, 1, 0], dtype=torch.float)
)
self.gamma = gamma

def configure_optimizers(self):
"""Configure the optimizers."""
return torch.optim.Adam(
[
{"params": self.encoder.parameters(), "lr": self.lr},
{"params": self.regressor.parameters(), "lr": self.lr},
]
)

def forward(self, x):
"""Perform the forward pass."""
xs = self.encoder(x).view(-1, self.encoder_output_dim)
theta = self.regressor(xs).view(-1, 2, 3)

grid = F.affine_grid(theta, x.size(), align_corners=False)
x_hat = F.grid_sample(x, grid, padding_mode="border", align_corners=False)

return x_hat, theta

@torch.no_grad()
def classify(self, x: torch.Tensor, split_size: int = 256):
"""Classify the input."""
x_hat, *_ = self(x)
x_hat = x_hat.unsqueeze(1).detach()
buffer = torch.stack([torch.from_numpy(img) for img in self._buffer]).to(
self.device
)
buffer = buffer.unsqueeze(0).detach()

losses = [] # classify in chunks to avoid OOM
for chunk in torch.split(buffer, split_size, dim=1):
chunk = chunk.repeat(x_hat.shape[0], 1, 1, 1, 1)
loss = F.mse_loss(
x_hat.repeat(1, chunk.shape[1], 1, 1, 1), chunk, reduction="none"
).mean(dim=(2, 3, 4))
losses.append(loss)
return torch.cat(losses, dim=1).argmin(dim=1)

def training_step(self, batch, batch_idx):
"""Perform a training step."""
loss = self._step(batch)
self.log_dict({f"{k}_train": v for k, v in loss.items() if k != "loss"})
return loss

def validation_step(self, batch, batch_idx):
"""Perform a validation step."""
loss = self._step(batch)
self.log_dict({f"{k}_val": v for k, v in loss.items() if k != "loss"})
return loss

def test_step(self, batch, batch_idx):
"""Perform a test step."""
loss = self._step(batch)
self.log_dict({f"{k}_test": v for k, v in loss.items() if k != "loss"})
return loss

def _step(self, batch):
"""Perform a training or validation step."""
x, y = batch
x_hat, theta_hat = self.forward(x)
exemplars = torch.stack(
[torch.from_numpy(self._buffer[i]) for i in y.shape_id]
).to(self.device)
theta = self.convert_parameters_to_matrix(y)
regression_loss = F.mse_loss(theta, theta_hat)
reconstruction_loss = F.mse_loss(exemplars, x_hat)
accuracy = (y.shape_id == self.classify(x)).float().mean()
return {
"accuracy": accuracy.item(),
"reconstruction_loss": reconstruction_loss.item(),
"regression_loss": regression_loss.item(),
"loss": self.gamma * regression_loss
+ (1 - self.gamma) * reconstruction_loss,
}
```
And the following visualisation callback:
```python
class VisualizationCallback(Callback):
"""Callback for visualizing reconstruction and classification."""

def __init__(self, canonical_images):
super().__init__()
self._canonical_images = canonical_images

def on_test_epoch_end(
self, trainer: pl.Trainer, pl_module: pl.LightningModule
) -> None:
"""Show the exemplars and the corresponding reconstructions."""
self.log_reconstructions(
pl_module,
self._canonical_images,
name="reconstructions_canonical",
)

def on_test_batch_end(
self,
trainer: pl.Trainer,
pl_module: pl.LightningModule,
outputs: Optional[STEP_OUTPUT],
batch: Any,
batch_idx: int,
dataloader_idx: int = 0,
) -> None:
if batch_idx == 0:
self.log_classification(
pl_module,
batch,
name="classification",
)

@staticmethod
@torch.no_grad()
def log_reconstructions(pl_module, x, name):
"""Log images and reconstructions"""
x = np.stack(x)
x_hat, *_ = pl_module(torch.from_numpy(x).to(pl_module.device))
images = draw_batch_and_reconstructions(x, to_numpy(x_hat))
pl_module.logger.log_image(name, images=[images])

@staticmethod
@torch.no_grad()
def log_classification(pl_module, batch, name):
x, y = batch
x, y = x.to(pl_module.device), y.to(pl_module.device)
x_hat, *_ = pl_module(x)
labels = pl_module.classify(x)
actual = np.stack([pl_module._buffer[i] for i in y.shape_id])
closest = np.stack([pl_module._buffer[i] for i in labels])
images = draw_batch_and_reconstructions(
to_numpy(x),
to_numpy(x_hat),
actual,
closest,
)
pl_module.logger.log_image(name, images=[images])
```

I'm training the model in a continual learning setting:
```python
for train_loader, val_loader, test_loader in zip(loaders):
trainer.fit(model, train_loader, val_loader)
trainer.fit_loop.max_epochs += cfg.trainer.max_epochs
trainer.test(model, test_loader)
```

I'm getting wildly different results with and without `@torch.no_grad` decorators in the callback and the `classify` method:

| without `@torch.no_grad` | with `@torch.no_grad` |
|---|---|
|  | |
|  | |

It seems like the test gradients are being leaked through the model `classify` method, even though I'm not even calculating the loss and calling `loss.backward()` anywhere. Shouldn't the gradients be cleared when starting `trainer.fit(loader)`? This behaviour was definitely unexpected to me and very hard to catch (my results looked too good to be true so I got suspicious).

### What version are you seeing the problem on?

v2.0

### How to reproduce the bug

Unfortunately I don't have a minimal reproducible example. I suspect the problem might arise from hacking the trainer to work in a continual learning setup?

### Error messages and logs

```
# Error messages and logs here please
```

### Environment

Current 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_

Start with SpatialTransformer.classify and the VisualizationCallback methods log_reconstructions and log_classification, then trace the repeated trainer.fit and trainer.test loop. Reproduce the difference with and without torch.no_grad and determine whether it comes from callback execution, model state, or the custom continual-learning loop; done means the cause is isolated and documented with a reproducible example.

Written by the indexing model from the issue text.