def test_memory_consumption_validation(tmp_path):
    """Test that the training batch is no longer in GPU memory when running validation.

    Cannot run with MPS, since there we can only measure shared memory and not dedicated, which device has how much
    memory allocated.

    """

    def get_memory():
        _clear_cuda_memory()
        return torch.cuda.memory_allocated(0)

    initial_memory = get_memory()

    class BoringLargeBatchModel(BoringModel):
        @property
        def num_params(self):
            return sum(p.numel() for p in self.parameters())

        def train_dataloader(self):
            # batch target memory >= 100x boring_model size
            batch_size = self.num_params * 100 // 32 + 1
            return DataLoader(RandomDataset(32, 5000), batch_size=batch_size)

        def val_dataloader(self):
            return self.train_dataloader()

        def training_step(self, batch, batch_idx):
            # there is a batch and the boring model, but not two batches on gpu, assume 32 bit = 4 bytes
            lower = 101 * self.num_params * 4
            upper = 201 * self.num_params * 4
            current = get_memory()
            assert lower < current
            assert current - initial_memory < upper
            return super().training_step(batch, batch_idx)

        def validation_step(self, batch, batch_idx):
            # there is a batch and the boring model, but not two batches on gpu, assume 32 bit = 4 bytes
            lower = 101 * self.num_params * 4
            upper = 201 * self.num_params * 4
            current = get_memory()
            assert lower < current
            assert current - initial_memory < upper
            return super().validation_step(batch, batch_idx)

    _clear_cuda_memory()
    trainer = Trainer(
        accelerator="gpu",
        devices=1,
        default_root_dir=tmp_path,
        fast_dev_run=2,
        enable_model_summary=False,
    )
    trainer.fit(BoringLargeBatchModel())
Test that the training batch is no longer in GPU memory when running validation.

    Cannot run with MPS, since there we can only measure shared memory and not dedicated, which device has how much
    memory allocated.

    
tests/tests_pytorch/loops/test_evaluation_loop.py