def test_train_step_no_return(tmp_path):
    """Tests that only training_step raises a warning when nothing is returned in case of automatic_optimization."""

    class TestModel(BoringModel):
        def training_step(self, batch):
            self.training_step_called = True
            loss = self.step(batch[0])
            self.log("a", loss, on_step=True, on_epoch=True)

        def validation_step(self, batch, batch_idx):
            self.validation_step_called = True

    model = TestModel()
    trainer_args = {"default_root_dir": tmp_path, "fast_dev_run": 2}
    trainer = Trainer(**trainer_args)

    Closure.warning_cache.clear()

    with pytest.warns(UserWarning, match=r"training_step` returned `None"):
        trainer.fit(model)

    assert model.training_step_called
    assert model.validation_step_called

    model = TestModel()
    model.automatic_optimization = False
    trainer = Trainer(**trainer_args)

    Closure.warning_cache.clear()

    with no_warning_call(UserWarning, match=r"training_step` returned `None"):
        trainer.fit(model)
Tests that only training_step raises a warning when nothing is returned in case of automatic_optimization.
tests/tests_pytorch/loops/test_training_loop_flow_scalar.py