import gc
import torch


def clear_memory():
    """Clean memory for `cuda` or `mps` device"""
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
        torch.cuda.synchronize()
    if torch.backends.mps.is_available():
        torch.mps.empty_cache()
        torch.mps.synchronize()
    gc.collect()
    print("Memory cleared\n")
