Update app.py

app.py

@@ -258,6 +258,15 @@ def load_dataset():
     
     return hf_dataset
 
+# Memory cleaning function
+def clean_memory():
+    gc.collect()
+    if torch.cuda.is_available():
+        for i in range(torch.cuda.device_count()):
+            with torch.cuda.device(f'cuda:{i}'):
+                torch.cuda.empty_cache()
+                torch.cuda.reset_peak_memory_stats()
+
 def train_model(progress=gr.Progress()):
     # Clean memory before starting
     clean_memory()
@@ -373,13 +382,4 @@ if __name__ == "__main__":
     
     # Create and launch the UI
     demo = create_ui()
-    demo.launch()
-
-# Memory cleaning function
-def clean_memory():
-    gc.collect()
-    if torch.cuda.is_available():
-        for i in range(torch.cuda.device_count()):
-            with torch.cuda.device(f'cuda:{i}'):
-                torch.cuda.empty_cache()
-                torch.cuda.reset_peak_memory_stats() 
+    demo.launch()