Spaces:
Sleeping
Sleeping
Update app.py
Browse files
app.py
CHANGED
@@ -91,7 +91,8 @@ def train_step(file=None, start_idx=0):
|
|
91 |
print("Exceeded GPU quota.")
|
92 |
if not os.path.exists('./checkpoint'):
|
93 |
os.makedirs('./checkpoint')
|
94 |
-
|
|
|
95 |
return start_idx # Trả về start_idx để lưu lại vị trí
|
96 |
|
97 |
start_idx = end_idx
|
|
|
91 |
print("Exceeded GPU quota.")
|
92 |
if not os.path.exists('./checkpoint'):
|
93 |
os.makedirs('./checkpoint')
|
94 |
+
print('Save checkpoint')
|
95 |
+
torch.save(model.state_dict(), "./checkpoint/model.pt")
|
96 |
return start_idx # Trả về start_idx để lưu lại vị trí
|
97 |
|
98 |
start_idx = end_idx
|