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

Code before:
            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}

        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()

    # test we have good test accuracy
    tutils.assert_ok_model_acc(trainer, key='test_loss', thr=0.5)


def test_running_test_no_val(tmpdir):
    """Verify `test()` works on a model with no `val_loader`."""
    class ModelTrainTest(BoringModel):
        def __init__(self, *args, **kwargs):
            super().__init__(*args, **kwargs)

        def val_loader(self):

Review comment (performance):
```suggestion
        def test_step(self, batch, batch_idx):
            output = self.layer(batch)
            loss = self.loss(batch, output)
            self.log('test_loss', loss)
            return {"y": loss}
```
same here.

Code after:
        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()

    # test we have good test accuracy
    tutils.assert_ok_model_acc(trainer, key='test_loss', thr=0.5)


def test_running_test_no_val(tmpdir):
    """Verify `test()` works on a model with no `val_dataloader`. It performs
    train and test only"""
    class ModelTrainTest(BoringModel):

        def val_dataloader(self):
            pass

        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 = ModelTrainTest()

    # 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=1,
        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()