code: def test_on_train_epoch_end_metrics_collection(tmpdir):
      """Test that progress bar metrics also get collected at the end of an epoch."""
      num_epochs = 3
  
      class CurrentModel(BoringModel):
          def training_step(self, *args, **kwargs):
              output = super().training_step(*args, **kwargs)
              self.log_dict(
                  {
                      "step_metric": torch.tensor(-1),
                      "shared_metric": 100,
                  },
                  logger=False,
                  prog_bar=True,
              )
              return output
  
          def on_train_epoch_end(self):
              epoch = self.current_epoch
              # both scalar tensors and Python numbers are accepted
              self.log_dict(
                  {
                      f"epoch_metric_{epoch}": torch.tensor(epoch),
                      "shared_metric": 111,
                  },
                  logger=False,
                  prog_bar=True,
              )
  
      model = CurrentModel()
      trainer = Trainer(
          max_epochs=num_epochs,
          default_root_dir=tmpdir,
          overfit_batches=2,
      )
      trainer.fit(model)
      assert (
          trainer.state.finished
      ), f"Training failed with {trainer.state}"
      metrics = trainer.progress_bar_callback.get_metrics(
          trainer, model
      )
  
      # metrics added in training step should be unchanged by epoch end method
      assert metrics["step_metric"] == -1
      # a metric shared in both methods gets overwritten by epoch_end
      assert metrics["shared_metric"] == 111
      # metrics are kept after each epoch
      for i in range(num_epochs):
          assert metrics[f"epoch_metric_{i}"] == i
  
imports: from functools import partial, update_wrapper
  from inspect import getmembers, isfunction
  from unittest import mock
  from unittest.mock import ANY, PropertyMock
  import pytest
  import torch
  from lightning.fabric.utilities.imports import _TORCH_GREATER_EQUAL_2_0
  from lightning.pytorch import Callback, LightningDataModule, LightningModule, Trainer, __version__
  from lightning.pytorch.demos.boring_classes import BoringDataModule, BoringModel, RandomDataset
  from lightning.pytorch.utilities.model_helpers import is_overridden
  from torch import Tensor
  from torch.utils.data import DataLoader
  from tests_pytorch.helpers.runif import RunIf
label_desc: Use image node_id 2 for calling a global function with example usage: test_on_train_epoch_end_metrics_collection(tmpdir) and returns: output