def test_train_step_no_return(tmp_path, single_cb: bool):
    """Tests that only training_step can be used."""

    class CB(Callback):
        def on_train_batch_end(self, trainer, pl_module, outputs, *_):
            assert "loss" in outputs

        def on_validation_batch_end(self, trainer, pl_module, outputs, *_):
            assert "x" in outputs

        def on_test_batch_end(self, trainer, pl_module, outputs, *_):
            assert "x" in outputs

    class TestModel(BoringModel):
        def on_train_batch_end(self, outputs, *_):
            assert "loss" in outputs

        def on_validation_batch_end(self, outputs, *_):
            assert "x" in outputs

        def on_test_batch_end(self, outputs, *_):
            assert "x" in outputs

    model = TestModel()

    trainer = Trainer(
        callbacks=CB() if single_cb else [CB()],
        default_root_dir=tmp_path,
        limit_train_batches=2,
        limit_val_batches=2,
        max_epochs=1,
        log_every_n_steps=1,
        enable_model_summary=False,
    )

    assert any(isinstance(c, CB) for c in trainer.callbacks)

    trainer.fit(model)
Tests that only training_step can be used.
tests/tests_pytorch/callbacks/test_callback_hook_outputs.py