Lightning-AI / Lightning-AI/pytorch-lightning
Bug: automatic logging doesn't log metric on steps if .update is used
@SkafteNicki is already working on this.
Since Sep 13, 2025.
- Dominant language
- Python
- Stars
- 31.4k
- Forks
- 3.8k
- Avg merge
- 6d 7h
- Merged PRs (30d)
- 6
Description
### Bug description
Before logging a metric inside a LightningModule it's state must be updated. There are two ways to do this: by calling `.update()` on the metric object, or by calling `.forward()` on it.
The issue is: specifically when using `.update`, and then passing the metric obj itself to the `self.log` command, the metric won't be logged except at the end of the epoch, even when the `on_step` kwarg is set to True. When using `forward`, the issue doesn't happen.
The root cause to this is that inside the logger, it checks the `self._forward_cache` attribute of the passed metric in order to obtain the value to log, but this attribute is only specifically set to a non-None value when the metric is invoked with `forward`, and not when it is invoked with `update`.
This issue has lead teams I'm working with to a lot of confusion - there's a tendency to prefer using `.update` because it is simpler and has less side effects than `forward`, and the metrics we work with don't need back propagation anyway. But subtle differences like this issue are really hard to understand without diving deep into the Lightning source code, and the user experience is likely to be much better if behavior was more consistent.
I suggest ensuring that, when a metric is passed to the logger, the value logged would be based on the current state of the logger, regardless of if the state was updated with forward or update.
Below I attach reproducing code and a description of how I found the source of the issue.
### What version are you seeing the problem on?
master
### How to reproduce the bug
```python
Toggle between the two different state update methods in the step function, and look at the output csv file to observe the difference in behavior - steps are only logged when calling self.accuracy(y_hat, y) and not when calling self.accuracy.update(y_hat, y).
from typing import Dict, Iterator
from lightning import LightningModule, LightningDataModule, Trainer
from lightning.pytorch.loggers import CSVLogger
import torch
from torch import nn
import torch.utils
from torch.utils.data import IterableDataset, DataLoader
from torchmetrics import Accuracy
import numpy as np
class RandomTensorDataset(IterableDataset):
def __iter__(self) -> Iterator:
while True:
yield np.random.random((1,)).astype(np.float32).squeeze(), 1.0
class MyDataModule(LightningDataModule):
def __init__(self):
super().__init__()
self._train_dataloader = DataLoader(RandomTensorDataset())
self._val_dataloader = DataLoader(RandomTensorDataset())
def train_dataloader(self):
return self._train_dataloader
def val_dataloader(self):
return self._val_dataloader
class MyLitClassifier(LightningModule):
def __init__(self):
super().__init__()
self.model = nn.Linear(1,1)
self.accuracy = Accuracy("binary")
self.loss_func = nn.CrossEntropyLoss()
def forward(self, x):
return self.model(x)
def step(self, batch: Dict[str, torch.Tensor]):
x, y = batch
y_hat = self(x)
loss = self.loss_func(y_hat, y)
# If we do this, accuracy will *not* be logged on_step
self.accuracy.update(y_hat, y)
# If we do this, accuracy *will* be logged on_step
# accuracy = self.accuracy(y_hat, y)
return loss
def training_step(self, batch: Dict[str, torch.Tensor], batch_idx: int):
loss = self.step(batch)
self.log(f"acc", self.accuracy, on_step=True, on_epoch=True)
return loss
def validation_step(self, batch: Dict[str, torch.Tensor], batch_idx: int):
loss = self.step(batch)
return loss
def configure_optimizers(self):
optimizer = torch.optim.SGD(
self.parameters(), lr=0.01, momentum=0.9
)
return optimizer
if __name__ == "__main__":
datamodule = MyDataModule()
model = MyLitClassifier()
logger = CSVLogger("logs")
trainer = Trainer(log_every_n_steps=5, limit_train_batches=25, limit_val_batches=10, max_epochs=10, logger=logger)
trainer.fit(
model=model,
datamodule=datamodule
)
```
### Error messages and logs
n/a
### Environment
Lightning: 2.3.3
Pytorch: 2.3.1
Python: 3.12.3
Machine: Ubuntu 18.04.4 LTS
### More info
About finding the root cause of the issue:
In the module lightning.pytorch.trainer.connectors.loger_connector.result.py, the _get_cache method of the ResultCollection object behaves differently for on_step vs. on_epoch - specifically in on_step, it obtains the cache by taking result_metric._forward_cache.
Looking at the _ResultMetric class (which wraps every metric sent for logging), it's forward method simply calls update, and update sets _ResultMetric._forward_cache to be equal to the _forward_cache value of the metric to be logged. However, in the torchmetrics.Metric class, only the forward method actually sets this value.
Thus, if forward was never called on the original metric, this value is None.
I'd be happy to provide more details if someone is looking at this bug and finds what I've written here useful but insufficient :)
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Assessment
This issue has not been assessed yet.