Update app.py

app.py

@@ -598,6 +598,8 @@ def run_train_epoch(num):
    for i in range(num):
     # model.train(True)
         yield pro.run_train_epoch()
+   save_pretrained(pro.dir_model)
+   return 'save model '
        
 @spaces.GPU
 def init_training():