From 9e2baa4f03ff3117e392faa04c71686b0e54803a Mon Sep 17 00:00:00 2001 From: Eren Date: Fri, 10 Aug 2018 17:49:21 +0200 Subject: [PATCH] Dont tim model param names since we don't use DataPatallel initilization --- utils/generic_utils.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/utils/generic_utils.py b/utils/generic_utils.py index c9a84fb6..c58efd75 100644 --- a/utils/generic_utils.py +++ b/utils/generic_utils.py @@ -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(),