Function free_memory

diffusers/src/diffusers/training_utils.py:263–272

Runs garbage collection. Then clears the cache of the available accelerator.

()

def free_memory():
    """Runs garbage collection. Then clears the cache of the available accelerator."""
    gc.collect()

    if torch.cuda.is_available():
        torch.cuda.empty_cache()
    elif torch.backends.mps.is_available():
        torch.mps.empty_cache()
    elif is_torch_npu_available():
        torch_npu.empty_cache()


# Adapted from torch-ema https://github.com/fadel/pytorch_ema/blob/master/torch_ema/ema.py#L14