wassname/TTS
1
0
Fork 0
mirror of https://github.com/wassname/TTS.git synced 2026-09-09 11:16:00 +08:00
Code Issues Packages Projects Releases Wiki Activity

Checkpoint stop token optimizer

Browse Source
This commit is contained in:
Eren G
2018-07-20 12:23:44 +02:00
parent fce6bd27b8
commit dfd0bc1831
2 changed files with 4 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
utils/generic_utils.py
+2 -1
View File
@@ -78,7 +78,7 @@ def _trim_model_state_dict(state_dict):
return new_state_dict
def save_checkpoint(model, optimizer, model_loss, out_path,
def save_checkpoint(model, optimizer, optimizer_st, model_loss, out_path,
current_step, epoch):
checkpoint_path = 'checkpoint_{}.pth.tar'.format(current_step)
checkpoint_path = os.path.join(out_path, checkpoint_path)
@@ -87,6 +87,7 @@ def save_checkpoint(model, optimizer, model_loss, out_path,
new_state_dict = _trim_model_state_dict(model.state_dict())
state = {'model': new_state_dict,
'optimizer': optimizer.state_dict(),
'optimizer_st': optimizer_st.state_dict(),
'step': current_step,
'epoch': epoch,
'linear_loss': model_loss,