Update app.py

app.py

@@ -595,7 +595,7 @@ pro=TrinerModelVITS(dir_model=dir_model,
 @spaces.GPU
 def run_train_epoch(num):
    pro.init_training()
-   for i in range(10):
+   for i in range(num):
     # model.train(True)
         yield pro.run_train_epoch()
        
@@ -619,7 +619,8 @@ def init_wandb():
 @spaces.GPU
 def save_pretrained(path):
     pro.save_pretrained(path)
-    pro.init_training()
+    
+    pro.load_model()
     return 'save_pretrained'
 
 with gr.Blocks() as interface: