Repository: Lightning-AI/pytorch-lightning
File: tests/models/test_cpu.py (Python)

Code before:
        default_root_dir=tmpdir,
        max_epochs=1,
        progress_bar_refresh_rate=0,
        weights_summary="top",
        limit_train_batches=0.2,
        limit_val_batches=0.2,
    )

    model = ModelSpecifiedOptimizer(optimizer_name="LBFGS", learning_rate=0.004)
    tpipes.run_model_test_without_loggers(trainer_options, model, min_acc=0.25)


def test_default_logger_callbacks_cpu_model(tmpdir):
    """Test each of the trainer options."""
    trainer_options = dict(
        default_root_dir=tmpdir,
        max_epochs=1,
        gradient_clip_val=1.0,
        overfit_batches=0.20,
        progress_bar_refresh_rate=0,
        limit_train_batches=0.01,
        limit_val_batches=0.01,
    )

    model = BoringModel()
    tpipes.run_model_test_without_loggers(trainer_options, model, min_acc=0.01)

    # test freeze on cpu
    model.freeze()
    model.unfreeze()


def test_running_test_after_fitting(tmpdir):
    """Verify test() on fitted model."""
    class ModelTrainValTest(BoringModel):
        def __init__(self, *args, **kwargs):
            super().__init__(*args, **kwargs)

        def validation_step(self, batch, batch_idx):
            output = self.layer(batch)
            loss = self.loss(batch, output)
            return {"x": loss}

        def validation_epoch_end(self, outputs) -> None:
            val_loss = torch.stack([x["x"] for x in outputs]).mean()
            self.log('val_loss', val_loss)

        def test_step(self, batch, batch_idx):
            output = self.layer(batch)
            loss = self.loss(batch, output)
            return {"y": loss}

Review comment (performance):
```suggestion
        def validation_step(self, batch, batch_idx):
            output = self.layer(batch)
            loss = self.loss(batch, output)
            self.log('val_loss', loss)

        def test_step(self, batch, batch_idx):
            output = self.layer(batch)
            loss = self.loss(batch, output)
            self.log('test_loss', loss)
```
let lightning aggregate the loss over the batches.

Code after:
        gradient_clip_val=1.0,
        overfit_batches=0.20,
        progress_bar_refresh_rate=0,
        limit_train_batches=0.01,
        limit_val_batches=0.01,
    )

    model = BoringModel()
    tpipes.run_model_test_without_loggers(trainer_options, model, min_acc=0.01)

    # test freeze on cpu
    model.freeze()
    model.unfreeze()


def test_running_test_after_fitting(tmpdir):
    """Verify test() on fitted model."""
    class ModelTrainValTest(BoringModel):

        def validation_epoch_end(self, outputs) -> None:
            val_loss = torch.stack([x["x"] for x in outputs]).mean()
            self.log('val_loss', val_loss)

        def test_epoch_end(self, outputs) -> None:
            test_loss = torch.stack([x["y"] for x in outputs]).mean()
            self.log('test_loss', test_loss)

    model = ModelTrainValTest()

    # logger file to get meta
    logger = tutils.get_default_logger(tmpdir)

    # logger file to get weights
    checkpoint = tutils.init_checkpoint_callback(logger)

    # fit model
    trainer = Trainer(
        default_root_dir=tmpdir,
        progress_bar_refresh_rate=0,
        max_epochs=2,
        limit_train_batches=0.4,
        limit_val_batches=0.2,
        limit_test_batches=0.2,
        callbacks=[checkpoint],
        logger=logger,
    )
    result = trainer.fit(model)

    assert result == 1, 'training failed to complete'

    trainer.test()