Fix bug at continue training.

yawudede · Nov 9, 2018 · bcbb699 · bcbb699
1 parent 47ad060
commit bcbb699
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/models/base_model.py b/models/base_model.py
@@ -39,11 +39,14 @@ def setup(self, opt, parser=None):
         if self.isTrain:
             self.schedulers = [networks.get_scheduler(optimizer, opt) for optimizer in self.optimizers]
 
-        if not self.isTrain or opt.continue_train:
+        if not self.isTrain:
             if opt.load_by_iter: # load network by iteration.
                 self.load_network('iter_' + opt.iter)
             else:
                 self.load_networks(opt.epoch)
+        else: # When continue training, only supported by loading from some epoch, not some iteration
+            if opt.continue_train:
+                self.load_networks(opt.epoch)
         self.print_networks(opt.verbose)
 
     # make models eval mode during test time