Update app.py

app.py

@@ -51,7 +51,10 @@ def train_batch(dataloader):
 
         elapsed_time = time.time() - start_time
         if elapsed_time > 1:  # Dừng trước 59 giây để đảm bảo không vượt hạn ngạch
+            if not os.path.exists('./checkpoint'):
+                os.makedirs('./checkpoint')
             torch.save(model.state_dict(), "./checkpoint/model.pt")
+            
             return False, "Checkpoint saved. Training paused."
 
     return True, "Batch training completed."
@@ -85,7 +88,9 @@ def train_step(file=None):
 
         start_idx = end_idx
         time.sleep(2)  # Nghỉ 2 giây giữa các phiên huấn luyện
-
+        
+    if not os.path.exists('./checkpoint'):
+        os.makedirs('./checkpoint')
     torch.save(model.state_dict(), "./checkpoint/model.pt")
     return "Training completed and model saved."