拆分utils函数

Leroll · Oct 27, 2023 · 2fabeff · 2fabeff
1 parent 418774e
commit 2fabeff
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 14 deletions.
diff --git a/utils/clean_memory.py b/utils/clean_memory.py
@@ -0,0 +1,10 @@
+#---------------------------------
+#  Function to clean RAM & vRAM
+#--------------------------------- 
+import gc
+import ctypes
+import torch
+def clean_memory():
+    gc.collect()
+    ctypes.CDLL("libc.so.6").malloc_trim(0)
+    torch.cuda.empty_cache()
diff --git a/utils/load_google_drive.py b/utils/load_google_drive.py
@@ -1,5 +1,3 @@
-
-
 #---------------------------------
 #  colab 加载 google drive
 #--------------------------------- 
@@ -16,15 +14,3 @@ def load_google_drive(path='/'):
 
     os.chdir(root_path + path)
     print(f"cwd: {os.getcwd()}")
-
-
-#---------------------------------
-#  Function to clean RAM & vRAM
-#--------------------------------- 
-import gc
-import ctypes
-import torch
-def clean_memory():
-    gc.collect()
-    ctypes.CDLL("libc.so.6").malloc_trim(0)
-    torch.cuda.empty_cache()