def test_logging_in_callbacks_with_log_function(tmp_path):
    """Tests ensure self.log can be used directly in callbacks."""

    class LoggingCallback(callbacks.Callback):
        def on_train_start(self, trainer, pl_module):
            self.log("on_train_start", 1)

        def on_train_epoch_start(self, trainer, pl_module):
            self.log("on_train_epoch_start", 2)

        def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx):
            self.log("on_train_batch_end", 3)

        def on_train_epoch_end(self, trainer, pl_module):
            self.log("on_train_epoch_end", 5)

    model = BoringModel()
    trainer = Trainer(
        default_root_dir=tmp_path,
        limit_train_batches=1,
        limit_val_batches=1,
        max_epochs=1,
        enable_model_summary=False,
        callbacks=[LoggingCallback()],
    )
    trainer.fit(model)

    expected = {
        "on_train_start": 1,
        "on_train_epoch_start": 2,
        "on_train_batch_end": 3,
        "on_train_epoch_end": 5,
    }
    assert trainer.callback_metrics == expected
Tests ensure self.log can be used directly in callbacks.
tests/tests_pytorch/trainer/logging_/test_train_loop_logging.py