Update train_second.py
Browse files- train_second.py +1 -1
train_second.py
CHANGED
@@ -585,7 +585,7 @@ def main(config_path):
|
|
585 |
'iters': iters,
|
586 |
'epoch': epoch,
|
587 |
}
|
588 |
-
save_path = osp.join(
|
589 |
torch.save(state, save_path)
|
590 |
|
591 |
loss_test = 0
|
|
|
585 |
'iters': iters,
|
586 |
'epoch': epoch,
|
587 |
}
|
588 |
+
save_path = osp.join(log_dir, f'2nd_STAGE_{epoch*len(train_dataloader)+i}.pth')
|
589 |
torch.save(state, save_path)
|
590 |
|
591 |
loss_test = 0
|