def test_multiple_trainer_constant_memory_allocated(tmp_path):
    """This tests ensures calling the trainer several times reset the memory back to 0."""

    class TestModel(BoringModel):
        def training_step(self, batch, batch_idx):
            loss = super().training_step(batch, batch_idx)
            self.log("train_loss", loss["loss"])
            return loss

        def configure_optimizers(self):
            return torch.optim.Adam(self.layer.parameters(), lr=0.1)

    class Check(Callback):
        def on_train_epoch_start(self, trainer, *_):
            assert isinstance(trainer.strategy.model, DistributedDataParallel)

    def current_memory():
        # before measuring the memory force release any leftover allocations, including CUDA tensors
        gc.collect()
        return torch.cuda.memory_allocated(0)

    initial = current_memory()

    model = TestModel()
    trainer_kwargs = {
        "default_root_dir": tmp_path,
        "fast_dev_run": True,
        "accelerator": "gpu",
        "devices": 1,
        "strategy": "ddp",
        "enable_progress_bar": False,
        "callbacks": Check(),
    }
    trainer = Trainer(**trainer_kwargs)
    trainer.fit(model)

    assert trainer.strategy.model is model
    assert list(trainer.optimizers[0].state.values())[0]["exp_avg_sq"].device == torch.device("cpu")
    assert trainer.callback_metrics["train_loss"].device == torch.device("cpu")

    assert current_memory() <= initial

    deepcopy(trainer)

    assert current_memory() <= initial

    trainer_2 = Trainer(**trainer_kwargs)
    trainer_2.fit(model)

    assert current_memory() <= initial
This tests ensures calling the trainer several times reset the memory back to 0.
tests/tests_pytorch/trainer/test_trainer.py