Update app.py

app.py

@@ -90,7 +90,7 @@ def train_step(file=None, start_idx=0):
         try:
             success, message = train_batch(dataloader)
             if not success:
-                return start_idx  # Trả về start_idx nếu lỗi xảy ra
+                return start_idx, "./checkpoint/model.pt  # Trả về start_idx nếu lỗi xảy ra
 
         except HTMLError as e:
             print(e)
@@ -98,20 +98,20 @@ def train_step(file=None, start_idx=0):
                 os.makedirs('./checkpoint')
             print('Save checkpoint')
             torch.save(model.state_dict(), "./checkpoint/model.pt")
-            return start_idx  # Trả về start_idx để lưu lại vị trí
+            return start_idx, "./checkpoint/model.pt"  # Trả về start_idx để lưu lại vị trí
 
         start_idx = end_idx
 
     if not os.path.exists('./checkpoint'):
         os.makedirs('./checkpoint')
     torch.save(model.state_dict(), "./checkpoint/model.pt")
-    return start_idx
+    return start_idx, "./checkpoint/model.pt
 
 
 if __name__ == "__main__":
     iface = gr.Interface(
         fn=train_step,
         inputs=[gr.File(label="Upload CSV"), gr.Textbox()],
-        outputs="text"
+        outputs=["text", gradio.File()]
     )
     iface.launch()