code: def test_metrics_reset(tmpdir):
      """Tests that metrics are reset correctly after the end of the train/val/test epoch.
  
      Taken from: `Metric Test for Reset`_
  
      """
  
      class TestModel(BoringModel):
          def __init__(self) -> None:
              super().__init__()
              self.layer = torch.nn.Linear(32, 1)
  
              for stage in ["train", "val", "test"]:
                  acc = BinaryAccuracy()
                  acc.reset = mock.Mock(side_effect=acc.reset)
                  ap = BinaryAveragePrecision()
                  ap.reset = mock.Mock(side_effect=ap.reset)
                  self.add_module(f"acc_{stage}", acc)
                  self.add_module(f"ap_{stage}", ap)
  
          def forward(self, x):
              return self.layer(x)
  
          def _step(self, stage, batch):
              labels = (
                  batch.detach().sum(1) > 0
              ).float()  # Fake some targets
              logits = self.forward(batch)
              loss = (
                  torch.nn.functional.binary_cross_entropy_with_logits(
                      logits, labels.unsqueeze(1)
                  )
              )
              probs = torch.sigmoid(logits.detach())
              self.log(f"loss/{stage}", loss)
  
              acc = self._modules[f"acc_{stage}"]
              ap = self._modules[f"ap_{stage}"]
  
              labels_int = labels.to(torch.long)
              acc(probs.flatten(), labels_int)
              ap(probs.flatten(), labels_int)
  
              # Metric.forward calls reset so reset the mocks here
              acc.reset.reset_mock()
              ap.reset.reset_mock()
  
              self.log(f"{stage}/accuracy", acc)
              self.log(f"{stage}/ap", ap)
  
              return loss
  
          def training_step(self, batch, batch_idx):
              return self._step("train", batch)
  
          def validation_step(self, batch, batch_idx):
              return self._step("val", batch)
  
          def test_step(self, batch, batch_idx):
              return self._step("test", batch)
  
          def _assert_epoch_end(self, stage):
              acc = self._modules[f"acc_{stage}"]
              ap = self._modules[f"ap_{stage}"]
  
              acc.reset.asset_not_called()
              ap.reset.assert_not_called()
  
          def on_train_epoch_end(self):
              self._assert_epoch_end("train")
  
          def on_validation_epoch_end(self):
              self._assert_epoch_end("val")
  
          def on_test_epoch_end(self):
              self._assert_epoch_end("test")
  
      def _assert_called(model, stage):
          acc = model._modules[f"acc_{stage}"]
          ap = model._modules[f"ap_{stage}"]
  
          acc.reset.assert_called_once()
          acc.reset.reset_mock()
  
          ap.reset.assert_called_once()
          ap.reset.reset_mock()
  
      model = TestModel()
      trainer = Trainer(
          default_root_dir=tmpdir,
          limit_train_batches=2,
          limit_val_batches=2,
          limit_test_batches=2,
          max_epochs=1,
      )
  
      trainer.fit(model)
      _assert_called(model, "train")
      _assert_called(model, "val")
  
      trainer.validate(model)
      _assert_called(model, "val")
  
      trainer.test(model)
      _assert_called(model, "test")
  
imports: from unittest import mock
  import torch
  from lightning_utilities import module_available
  from torch import tensor
  from torch.nn import Linear
  from torchmetrics import MetricCollection
  from torchmetrics.aggregation import SumMetric
  from torchmetrics.classification import BinaryAccuracy, BinaryAveragePrecision
  from torchmetrics.regression import MeanSquaredError
  from torchmetrics.wrappers import MultitaskWrapper
  from integrations.helpers import no_warning_call
  from integrations.lightning.boring_model import BoringModel
label_desc: Use image node_id 2 for calling a global function with example usage: test_metrics_reset(tmpdir) and returns: self, loss, self, self, self