def test__training_step__epoch_end__flow_scalar(tmp_path):
    """Tests that only training_step can be used."""

    class TestModel(DeterministicModel):
        def training_step(self, batch, batch_idx):
            acc = self.step(batch, batch_idx)
            acc = acc + batch_idx

            self.training_step_called = True
            return acc

        def backward(self, loss):
            return LightningModule.backward(self, loss)

    model = TestModel()
    model.val_dataloader = None

    trainer = Trainer(
        default_root_dir=tmp_path,
        limit_train_batches=2,
        limit_val_batches=2,
        max_epochs=2,
        log_every_n_steps=1,
        enable_model_summary=False,
    )
    trainer.fit(model)

    # make sure correct steps were called
    assert model.training_step_called

    # assert epoch end metrics were added
    assert len(trainer.callback_metrics) == 0
    assert len(trainer.progress_bar_metrics) == 0

    trainer.state.stage = RunningStage.TRAINING
    # make sure training outputs what is expected
    kwargs = {"batch": next(iter(model.train_dataloader())), "batch_idx": 0}
    train_step_out = trainer.fit_loop.epoch_loop.automatic_optimization.run(trainer.optimizers[0], 0, kwargs)

    assert isinstance(train_step_out["loss"], Tensor)
    assert train_step_out["loss"].item() == 171

    # make sure the optimizer closure returns the correct things
    opt_closure = trainer.fit_loop.epoch_loop.automatic_optimization._make_closure(kwargs, trainer.optimizers[0], 0)
    opt_closure_result = opt_closure()
    assert opt_closure_result.item() == 171
Tests that only training_step can be used.
tests/tests_pytorch/loops/test_training_loop_flow_scalar.py