Update app.py

app.py

@@ -70,6 +70,11 @@ def train_step(file=None, start_idx=0):
     if os.path.exists("./checkpoint/model.pt"):
         print("Loading checkpoint...")
         model.load_state_dict(torch.load("./checkpoint/model.pt"))
+    else:
+        print("Checkpoint not found, starting fresh...")
+        if not os.path.exists('./checkpoint'):
+            os.makedirs('./checkpoint')
+        torch.save(model.state_dict(), "./checkpoint/model.pt")
 
     batch_size = 8
     total_samples = len(global_data)