Fix: utils error in michi finetuning script (#218)

training/mochi-1/text_to_video_lora.py

@@ -40,10 +40,6 @@
 from dataset_simple import LatentEmbedDataset
 
 import sys
-
-
-sys.path.append("..")
-
 from utils import print_memory, reset_memory  # isort:skip
 
 

training/mochi-1/utils.py

@@ -0,0 +1,22 @@
+import gc
+import inspect
+from typing import Optional, Tuple, Union
+
+import torch
+
+logger = get_logger(__name__)
+
+def reset_memory(device: Union[str, torch.device]) -> None:
+    gc.collect()
+    torch.cuda.empty_cache()
+    torch.cuda.reset_peak_memory_stats(device)
+    torch.cuda.reset_accumulated_memory_stats(device)
+
+
+def print_memory(device: Union[str, torch.device]) -> None:
+    memory_allocated = torch.cuda.memory_allocated(device) / 1024**3
+    max_memory_allocated = torch.cuda.max_memory_allocated(device) / 1024**3
+    max_memory_reserved = torch.cuda.max_memory_reserved(device) / 1024**3
+    print(f"{memory_allocated=:.3f} GB")
+    print(f"{max_memory_allocated=:.3f} GB")
+    print(f"{max_memory_reserved=:.3f} GB")