Update app.py

app.py

@@ -34,7 +34,7 @@ def get_dataloader(start, end, batch_size=8):
     subset = torch.utils.data.Subset(global_data, range(start, end))
     return DataLoader(subset, batch_size=batch_size)
 
-@spaces.GPU(duration=20)
+@spaces.GPU(duration=120)
 def train_batch(dataloader):
     model.train()
     start_time = time.time()
@@ -88,8 +88,10 @@ def train_step(file=None, start_idx=0):
                 return start_idx  # Trả về start_idx nếu lỗi xảy ra
 
         except HTMLError as e:
-            print("Exceeded GPU quota, retrying in 10 seconds...")
-            time.sleep(10)
+            print("Exceeded GPU quota.")
+            if not os.path.exists('./checkpoint'):
+                os.makedirs('./checkpoint')
+                torch.save(model.state_dict(), "./checkpoint/model.pt")
             return start_idx  # Trả về start_idx để lưu lại vị trí
 
         start_idx = end_idx