code: def test_metrics_reset(tmpdir):
      """Tests that metrics are reset correctly after the end of the train/val/test epoch."""
  
      class TestModel(LightningModule):
          def __init__(self):
              super().__init__()
              self.layer = torch.nn.Linear(32, 1)
  
          def _create_metrics(self):
              acc = (
                  Accuracy(task="binary") if _TM_GE_0_11 else Accuracy()
              )
              acc.reset = mock.Mock(side_effect=acc.reset)
              ap = (
                  AvgPre(task="binary")
                  if _TM_GE_0_11
                  else AvgPre(num_classes=1, pos_label=1)
              )
              ap.reset = mock.Mock(side_effect=ap.reset)
              return acc, ap
  
          def setup(self, stage):
              fn = stage.value
              if fn == "fit":
                  for stage in ("train", "validate"):
                      acc, ap = self._create_metrics()
                      self.add_module(f"acc_{fn}_{stage}", acc)
                      self.add_module(f"ap_{fn}_{stage}", ap)
              else:
                  acc, ap = self._create_metrics()
                  stage = self.trainer.state.stage.value
                  self.add_module(f"acc_{fn}_{stage}", acc)
                  self.add_module(f"ap_{fn}_{stage}", ap)
  
          def forward(self, x):
              return self.layer(x)
  
          def _step(self, batch):
              fn, stage = (
                  self.trainer.state.fn.value,
                  self.trainer.state.stage.value,
              )
  
              logits = self(batch)
              loss = logits.sum()
              self.log(f"loss/{fn}_{stage}", loss)
  
              acc = self._modules[f"acc_{fn}_{stage}"]
              ap = self._modules[f"ap_{fn}_{stage}"]
  
              preds = torch.rand(len(batch))  # Fake preds
              labels = torch.randint(0, 1, [len(batch)])  # Fake targets
              acc(preds, labels)
              ap(preds, labels)
  
              # Metric.forward calls reset so reset the mocks here
              acc.reset.reset_mock()
              ap.reset.reset_mock()
  
              self.log(f"acc/{fn}_{stage}", acc)
              self.log(f"ap/{fn}_{stage}", ap)
  
              return loss
  
          def training_step(self, batch, batch_idx, *args, **kwargs):
              return self._step(batch)
  
          def validation_step(self, batch, batch_idx, *args, **kwargs):
              if self.trainer.sanity_checking:
                  return None
              return self._step(batch)
  
          def test_step(self, batch, batch_idx, *args, **kwargs):
              return self._step(batch)
  
          def configure_optimizers(self):
              optimizer = torch.optim.SGD(
                  self.layer.parameters(), lr=0.1
              )
              lr_scheduler = torch.optim.lr_scheduler.StepLR(
                  optimizer, step_size=1
              )
              return [optimizer], [lr_scheduler]
  
          def train_dataloader(self):
              return DataLoader(RandomDataset(32, 64))
  
          def val_dataloader(self):
              return DataLoader(RandomDataset(32, 64))
  
          def test_dataloader(self):
              return DataLoader(RandomDataset(32, 64))
  
      def _assert_called(model, fn, stage):
          acc = model._modules[f"acc_{fn}_{stage}"]
          ap = model._modules[f"ap_{fn}_{stage}"]
          acc.reset.assert_called_once()
          ap.reset.assert_called_once()
  
      model = TestModel()
      trainer = Trainer(
          default_root_dir=tmpdir,
          limit_train_batches=2,
          limit_val_batches=2,
          limit_test_batches=2,
          max_epochs=1,
          enable_progress_bar=False,
          num_sanity_val_steps=2,
          enable_checkpointing=False,
      )
  
      trainer.fit(model)
      _assert_called(model, "fit", "train")
      _assert_called(model, "fit", "validate")
  
      trainer.validate(model)
      _assert_called(model, "validate", "validate")
  
      trainer.test(model)
      _assert_called(model, "test", "test")
  
imports: import operator
  from functools import partial
  from unittest import mock
  from unittest.mock import Mock
  import pytest
  import torch
  from lightning.pytorch import LightningModule
  from lightning.pytorch.callbacks.callback import Callback
  from lightning.pytorch.demos.boring_classes import BoringModel, RandomDataset
  from lightning.pytorch.loggers import CSVLogger
  from lightning.pytorch.trainer import Trainer
  from lightning.pytorch.trainer.connectors.logger_connector.fx_validator import _FxValidator
  from lightning.pytorch.trainer.connectors.logger_connector.result import _ResultCollection
  from lightning.pytorch.utilities.exceptions import MisconfigurationException
  from lightning.pytorch.utilities.imports import _TORCHMETRICS_GREATER_EQUAL_0_9_1
  from lightning.pytorch.utilities.imports import _TORCHMETRICS_GREATER_EQUAL_0_11
  from lightning_utilities.core.imports import compare_version
  from torch.utils.data import DataLoader
  from torchmetrics import Accuracy, MeanAbsoluteError, MeanSquaredError, MetricCollection
  from torchmetrics import AveragePrecision
  from tests_pytorch.models.test_hooks import get_members
label_desc: Use image node_id 4 for calling a global function with example usage: test_metrics_reset(tmpdir) and returns: acc,  ap, self, loss, self, self, self, list,  list, DataLoader, DataLoader, DataLoader, None