Update main_train.py

yangtutuaka · Aug 30, 2018 · 5222842 · 5222842
1 parent 683a871
commit 5222842
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/TrainingCodes/dncnn_pytorch/main_train.py b/TrainingCodes/dncnn_pytorch/main_train.py
@@ -133,7 +133,8 @@ def log(*args, **kwargs):
     initial_epoch = findLastCheckpoint(save_dir=save_dir)  # load the last model in matconvnet style
     if initial_epoch > 0:
         print('resuming by loading epoch %03d' % initial_epoch)
-        model.load_state_dict(torch.load(os.path.join(save_dir, 'model_%03d.pth' % initial_epoch)))
+        # model.load_state_dict(torch.load(os.path.join(save_dir, 'model_%03d.pth' % initial_epoch)))
+        model = torch.load(os.path.join(save_dir, 'model_%03d.pth' % initial_epoch))
     model.train()
     # criterion = nn.MSELoss(reduction = 'sum')  # PyTorch 0.4.1
     criterion = sum_squared_error()