Dont tim model param names since we don't use DataPatallel initilization

This commit is contained in:
Eren 2018-08-10 17:49:21 +02:00
parent 96e2e3c776
commit 9e2baa4f03
1 changed files with 2 additions and 2 deletions

View File

@ -89,7 +89,7 @@ def save_checkpoint(model, optimizer, optimizer_st, model_loss, out_path,
checkpoint_path = os.path.join(out_path, checkpoint_path)
print(" | | > Checkpoint saving : {}".format(checkpoint_path))
new_state_dict = _trim_model_state_dict(model.state_dict())
new_state_dict = model.state_dict()
state = {
'model': new_state_dict,
'optimizer': optimizer.state_dict(),
@ -105,7 +105,7 @@ def save_checkpoint(model, optimizer, optimizer_st, model_loss, out_path,
def save_best_model(model, optimizer, model_loss, best_loss, out_path,
current_step, epoch):
if model_loss < best_loss:
new_state_dict = _trim_model_state_dict(model.state_dict())
new_state_dict = model.state_dict()
state = {
'model': new_state_dict,
'optimizer': optimizer.state_dict(),