Refactor memory cleaning into a single function

train_textual_inversion_XTI.py

@@ -1,6 +1,5 @@
 import importlib
 import argparse
-import gc
 import math
 import os
 import toml
@@ -9,6 +8,7 @@ from multiprocessing import Value
 from tqdm import tqdm
 import torch
 
+from library.device_utils import clean_memory
 from library.ipex_interop import init_ipex
 
 init_ipex()
@@ -286,9 +286,7 @@ def train(args):
         with torch.no_grad():
             train_dataset_group.cache_latents(vae, args.vae_batch_size, args.cache_latents_to_disk, accelerator.is_main_process)
         vae.to("cpu")
-        if torch.cuda.is_available():
-            torch.cuda.empty_cache()
-        gc.collect()
+        clean_memory()
 
         accelerator.wait_for_everyone()